def elementwise_binary_union(linear, use_sparsify):
    """Decorator to add support for `JAXSparse` tensors (including `BCOO`) to an
    element-wise binary operator such that the indices present in the result are
    are the union of the indices in the two operand.

    The primary use case for this is the `add` and `subtract` operators.

    There are requirements on the operator for this decorator to work correctly:

    - The operator must be element-wise.
    - The operator must be binary (two input tensors and one output tensor).
    - Both inputs must be of the same shape or one input must be a scalar.
    - The output must be of the same shape as the (non scalar) inputs.
    - The indices of the output must be the union of the indices of the inputs.
      This implies that func(0, 0) must be 0. As a result, if one operand is
      dense or a scalar, then the result will be dense.

    Additional arguments to the function (besides the input tensors) are not
    supported.

    Note that if the result of the operation is zero at some indices, including
    because the operands were zero at these indices, the zeros and indices are
    preserved.

    The `BCOO` format is the only supported one in all cases. Other formats are
    not supported when `use_sparsify` is `False`.

    Args:
        use_sparsify: indicates that the JAX `sparsify` transform supports this
            operation.
        linear: if `True`, mean that the operation is such that
            `op(a + b, c) == op(a, c) + op(b, c)` and
            `op(a, c + d) == op(a, c) + op(a, d)`.
    Returns:
        Wrapped function that supports `JAXSparse`.
    """

    def wrap_elementwise_binary_union(func):
        sparse_func = jax_sparse.sparsify(func) if use_sparsify else None

        @functools.wraps(func)
        def sparse_wrapper(x1, x2):
            if isinstance(x1, jax_sparse.JAXSparse):
                if isinstance(x2, jax_sparse.JAXSparse):
                    # x1 and x2 are sparse.
                    # The way we use `sparsify` it cannot know that the indices
                    # are the same, so we optimize this case here.
                    if (
                        x1.indices is x2.indices
                        and isinstance(x1, jax_sparse.BCOO)
                        and isinstance(x2, jax_sparse.BCOO)
                    ):
                        if not linear and not x1.unique_indices:
                            x1 = jax_sparse.bcoo_sum_duplicates(x1)
                            x2 = jax_sparse.bcoo_sum_duplicates(x2)
                        return jax_sparse.BCOO(
                            (func(x1.data, x2.data), x1.indices),
                            shape=x1.shape,
                            indices_sorted=x1.indices_sorted,
                            unique_indices=x1.unique_indices,
                        )
                    elif use_sparsify:
                        return sparse_func(x1, x2)
                    elif isinstance(x1, jax_sparse.BCOO) and isinstance(
                        x2, jax_sparse.BCOO
                    ):
                        x1 = bcoo_add_indices(x1, x2, sum_duplicates=not linear)
                        x2 = bcoo_add_indices(x2, x1, sum_duplicates=not linear)
                        return jax_sparse.BCOO(
                            (func(x1.data, x2.data), x1.indices),
                            shape=x1.shape,
                            indices_sorted=True,
                            unique_indices=True,
                        )
                    else:
                        ValueError(
                            "Unsupported sparse format: "
                            f"{x1.__class__} and {x2.__class__}"
                        )
                else:
                    # x1 is sparse, x2 is dense, densify x2.
                    x1 = x1.todense()
            elif isinstance(x2, jax_sparse.JAXSparse):
                # x1 is dense, x2 is sparse, densify x2.
                x2 = x2.todense()
            return func(x1, x2)

        return sparse_wrapper

    return wrap_elementwise_binary_union
Decorator to add support for `JAXSparse` tensors (including `BCOO`) to an
    element-wise binary operator such that the indices present in the result are
    are the union of the indices in the two operand.

    The primary use case for this is the `add` and `subtract` operators.

    There are requirements on the operator for this decorator to work correctly:

    - The operator must be element-wise.
    - The operator must be binary (two input tensors and one output tensor).
    - Both inputs must be of the same shape or one input must be a scalar.
    - The output must be of the same shape as the (non scalar) inputs.
    - The indices of the output must be the union of the indices of the inputs.
      This implies that func(0, 0) must be 0. As a result, if one operand is
      dense or a scalar, then the result will be dense.

    Additional arguments to the function (besides the input tensors) are not
    supported.

    Note that if the result of the operation is zero at some indices, including
    because the operands were zero at these indices, the zeros and indices are
    preserved.

    The `BCOO` format is the only supported one in all cases. Other formats are
    not supported when `use_sparsify` is `False`.

    Args:
        use_sparsify: indicates that the JAX `sparsify` transform supports this
            operation.
        linear: if `True`, mean that the operation is such that
            `op(a + b, c) == op(a, c) + op(b, c)` and
            `op(a, c + d) == op(a, c) + op(a, d)`.
    Returns:
        Wrapped function that supports `JAXSparse`.
    
keras/src/backend/jax/sparse.py