def elementwise_division(func):
    """Decorator to add support for `BCOO` sparse tensors to element-wise binary
    division and related operators.

    This decorator is designed for operations related to the division of two
    two operands (e.g. `divide`). It accepts `BCOO` tensors for both the
    dividend and the divisor, but handles them differently based on whether they
    are the dividend or the divisor.

    - If the divisor is sparse, it is densified and the result is dense because
      the result contains Inf or Nan outside of the indices of the dividend.
    - If the dividend is sparse and the divisor is dense, it finds occurrences
      of zeros and NaNs in the divisor. The result may therefore have more
      indices than there were in the dividend to return correct values where the
      divisor was zero or NaN.
    - If the dividend is sparse and the divisor is a scalar, it does the
      division element-wise. Note that the result is incorrectly sparse if the
      scalar divisor is zero.

    Args:
        func: The function to wrap.
    Returns:
        Wrapped function that supports `BCOO` sparse tensors.
    """
    sparse_func = jax_sparse.sparsify(func)

    @functools.wraps(func)
    def sparse_wrapper(x1, x2):
        if isinstance(x1, jax_sparse.JAXSparse):
            if isinstance(x2, jax_sparse.JAXSparse):
                # x1 is sparse and x2 is sparse.
                # Divisor is sparse, meaning we're doing divisions by zero
                # outside of x2.indices, so the result is dense. Densify both.
                x1 = x1.todense()
                x2 = x2.todense()
            elif isinstance(x1, jax_sparse.BCOO):
                if not hasattr(x2, "shape") or len(x2.shape) == 0:
                    # x1 is sparse BCOO, x2 is scalar, apply func element-wise.
                    return jax_sparse.BCOO(
                        (func(x1.data, x2), x1.indices),
                        shape=x1.shape,
                        indices_sorted=x1.indices_sorted,
                        unique_indices=x1.unique_indices,
                    )
                else:
                    # x1 is sparse BCOO, x2 is dense.
                    if not jax_utils.is_in_jax_tracing_scope(x2):
                        # Find zeros and nans in x2 and add indices to x1.
                        # 1. Create a dense mask for zeros and nans.
                        x2_zeros_and_nans = jnp.equal(x2, 0)
                        if not jnp.issubdtype(x2.dtype, jnp.integer):
                            x2_zeros_and_nans = jnp.logical_or(
                                x2_zeros_and_nans, jnp.isnan(x2)
                            )
                        # 2. Make it a BCOO of True values.
                        x2_zeros_and_nans = jax_sparse.bcoo_fromdense(
                            x2_zeros_and_nans,
                            n_batch=x1.n_batch,
                            n_dense=x1.n_dense,
                            index_dtype=x1.indices.dtype,
                        )
                        # 3. Add the indices to x1.
                        x1 = bcoo_add_indices(
                            x1, x2_zeros_and_nans, sum_duplicates=True
                        )
                    return sparse_func(x1, x2)
            else:
                raise ValueError(f"Unsupported sparse format: {x1.__class__}")
        elif isinstance(x2, jax_sparse.JAXSparse):
            # x1 is dense, x2 is sparse, densify x2
            x2 = x2.todense()
        return func(x1, x2)

    return sparse_wrapper
Decorator to add support for `BCOO` sparse tensors to element-wise binary
    division and related operators.

    This decorator is designed for operations related to the division of two
    two operands (e.g. `divide`). It accepts `BCOO` tensors for both the
    dividend and the divisor, but handles them differently based on whether they
    are the dividend or the divisor.

    - If the divisor is sparse, it is densified and the result is dense because
      the result contains Inf or Nan outside of the indices of the dividend.
    - If the dividend is sparse and the divisor is dense, it finds occurrences
      of zeros and NaNs in the divisor. The result may therefore have more
      indices than there were in the dividend to return correct values where the
      divisor was zero or NaN.
    - If the dividend is sparse and the divisor is a scalar, it does the
      division element-wise. Note that the result is incorrectly sparse if the
      scalar divisor is zero.

    Args:
        func: The function to wrap.
    Returns:
        Wrapped function that supports `BCOO` sparse tensors.
    
keras/src/backend/jax/sparse.py