def elementwise_unary(linear):
    """Decorator to add support for `BCOO` sparse tensors to a zero-preserving
    element-wise unary operator.

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

    - The operator must be element-wise
    - The operator must be unary (one input tensor and one output tensor)
    - The operator must return a tensor of the same shape, and if it is a
      `BCOO` tensor, the indices of the result must be the same. Therefore:
        - Reduction operations are not supported (e.g. `mean`).
        - Operations for which the result may be dense (e.g. `reciprocal`), or
          the sparse indices depend on the inputs are not supported (e.g.
          `clip`). This implies that `func(0)` must be 0.

    Additional arguments to the function (besides the input tensor) are
    supported as long as they cannot change the indices of the result. For
    instance,`round` is supported, but `clip` is not supported as
    `clip(x, 1.0, 2.0)` would always return a dense tensor.

    Note that if an input sparse tensor contains zero values, the indices and
    the zero values are preserved.

    Args:
        linear: if `True`, means that the operation is such that
            `op(a + b) == op(a) + op(b)`.
    Returns:
        Wrapped function that supports `BCOO` sparse tensors.
    """

    def wrap_elementwise_unary(func):
        @functools.wraps(func)
        def sparse_wrapper(x, *args, **kwargs):
            if isinstance(x, jax_sparse.BCOO):
                if not linear and not x.unique_indices:
                    x = jax_sparse.bcoo_sum_duplicates(x)
                return jax_sparse.BCOO(
                    (func(x.data, *args, **kwargs), x.indices), shape=x.shape
                )
            else:
                return func(x, *args, **kwargs)

        return sparse_wrapper

    return wrap_elementwise_unary
Decorator to add support for `BCOO` sparse tensors to a zero-preserving
    element-wise unary operator.

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

    - The operator must be element-wise
    - The operator must be unary (one input tensor and one output tensor)
    - The operator must return a tensor of the same shape, and if it is a
      `BCOO` tensor, the indices of the result must be the same. Therefore:
        - Reduction operations are not supported (e.g. `mean`).
        - Operations for which the result may be dense (e.g. `reciprocal`), or
          the sparse indices depend on the inputs are not supported (e.g.
          `clip`). This implies that `func(0)` must be 0.

    Additional arguments to the function (besides the input tensor) are
    supported as long as they cannot change the indices of the result. For
    instance,`round` is supported, but `clip` is not supported as
    `clip(x, 1.0, 2.0)` would always return a dense tensor.

    Note that if an input sparse tensor contains zero values, the indices and
    the zero values are preserved.

    Args:
        linear: if `True`, means that the operation is such that
            `op(a + b) == op(a) + op(b)`.
    Returns:
        Wrapped function that supports `BCOO` sparse tensors.
    
keras/src/backend/jax/sparse.py