A todense function that would work inside sparsified transforms

[edavisau]

Within sparsified-functions, does a todense type of function exist which works on both dense and sparse arrays? For a dense array this would just be the identity function, but for the sparse array it would call sparse.BCOO.todense().

In my use case:

@jit
@sparse.sparsify
def my_func(a: jnp.ndarray | sparse.BCOO) -> jnp.ndarray:
    # Section that works with or without sparse arrays 
    # ...
    # result: jnp.ndarray | sparse.BCOO = ...
    # ...
    
    result: jnp.ndarray = result.todense()
    # If result is:
    #  - dense -> return itself (no operation)
    #  - sparse -> convert to dense
    
    
    # Section that only works with dense arrays
    # ...
    # ...
    return result

I suppose it's like putting

if isinstance(result, sparse.BCOO):
    result = result.todense()

However, I don't think this works.

[jakevdp]

You can use sparse.todense for this

[edavisau]

I did not see this function, exactly what I needed, thank you!