keras/src/backend/jax/numpy.py/multiply
def multiply(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    if isinstance(x1, jax_sparse.BCOO):
        if isinstance(x2, jax_sparse.BCOO):
            # x1 is sparse, x2 is sparse.
            if x1.indices is x2.indices:
                # `bcoo_multiply_sparse` will not detect that the indices are
                # the same, optimize this case here.
                if not x1.unique_indices:
                    x1 = jax_sparse.bcoo_sum_duplicates(x1)
                    x2 = jax_sparse.bcoo_sum_duplicates(x2)
                return jax_sparse.BCOO(
                    (jnp.multiply(x1.data, x2.data), x1.indices),
                    shape=x1.shape,
                    indices_sorted=True,
                    unique_indices=True,
                )
            else:
                return jax_sparse.bcoo_multiply_sparse(x1, x2)
        else:
            # x1 is sparse, x2 is dense.
            out_data = jax_sparse.bcoo_multiply_dense(x1, x2)
            return jax_sparse.BCOO(
                (out_data, x1.indices),
                shape=x1.shape,
                indices_sorted=x1.indices_sorted,
                unique_indices=x1.unique_indices,
            )
    elif isinstance(x2, jax_sparse.BCOO):
        # x1 is dense, x2 is sparse.
        out_data = jax_sparse.bcoo_multiply_dense(x2, x1)
        return jax_sparse.BCOO(
            (out_data, x2.indices),
            shape=x2.shape,
            indices_sorted=x2.indices_sorted,
            unique_indices=x2.unique_indices,
        )
    return jnp.multiply(x1, x2)