def densifying_unary(func):
    """Decorator to add support for `JAXSparse` tensors (including `BCOO`) to a
    non-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.

    Additional arguments to the function (besides the input tensor) are
    supported. The returned result is a dense tensor.

    Args:
        func: The unary operator to wrap.
    Returns:
        Wrapped function that supports `JAXSparse` tensors.
    """

    @functools.wraps(func)
    def sparse_wrapper(x, *args, **kwargs):
        if isinstance(x, jax_sparse.JAXSparse):
            x = x.todense()
        return func(x, *args, **kwargs)

    return sparse_wrapper
Decorator to add support for `JAXSparse` tensors (including `BCOO`) to a
    non-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.

    Additional arguments to the function (besides the input tensor) are
    supported. The returned result is a dense tensor.

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