How to apply certain operation over sparse.BCOO array?

[edavisau]

Is there currently a way to do the following operation with sparse arrays?

import jax.numpy as jnp
from jax.experimental import sparse

probs = jnp.array(
    [
        [0.5, 0.0, 0.0, 0.6],
        [0.8, 0.0, 0.9, 0.1],
    ]
)
probs_sp = sparse.BCOO.fromdense(probs)

(1 - probs).prod(axis=1)
# returns [0.2, 0.018 ]

(1 - probs_sp).prod(axis=1)
# NotImplementedError: Subtraction between sparse and dense array.

For obvious reasons, (1 - probs_sp) is a bad idea. However, with .prod(...), we can actually compute this pretty efficiently with sparse arrays, you only need to perform these operations over the non-zero elements. The zeros do not affect the product.

Example: $(1 - 0.5) \times (1 - 0.6) = 0.2$ and $(1 - 0.8) \times (1 - 0.9) \times (1 - 0.1) = 0.018$

Thanks!

[edavisau]

I believe I could do
jnp.exp(sparse.sparsify(jnp.log1p)(-probs_sp).sum(axis=1).todense())

since
$$\prod_{i=1}^n 1 - p_i = \mathrm{exp} \left( \sum_{i=1}^n \log (1 + (-p_i)) \right)$$

but there must be a better way...

[jakevdp]

The reason this operation fails is because in general 1 - spmat is a fully-dense matrix, and JAX avoids implicit conversion to dense. If converting your matrix to dense is actually what you want to do, then you can do it this way:

(1 - probs_sp.todense()).prod(axis=1)

But I understand that for your specific computation, the sparse elements in the matrix become ones that do not contribute to the final product. JAX sparse does not have any support for this specific type of computation; you might be able to proceed by working in log-space where the ones become zeros and the multiplications become additions.

[edavisau]

Thanks, in the end the log-space solution seemed like the best way to proceed. And it is actually significantly faster for very sparse matrices as in my use case.

@jit
@sparse.sparsify
def func(probs):
    log_sum = jnp.log1p(-probs).sum(axis=1)
    return -jnp.expm1(log_sum)