Memory-efficient weighted sum over JAXSparse array

[quattro]

Hi all,

I've been working with some custom code that represents a dense centered/standardized matrix as a sparse linear operation, and have been running into some trouble trying to compute a memory efficient sum of squares.

Namely, what I'd like to compute (using an example with dense operations) would be jnp.sum((Z **2) * W), where Z is the centered/standardized matrix with shape (n,p), and W is a vector of weights with shape (n,).

If I represent Z as (G - C@B) * S, where C is a covariate matrix with shape (n,k), B is the mean-effect matrix with shape (k,p), and S is the scale-adjustment with shape (p,), this operation can be broken down into three terms:

term1 = jnp.einsum("np,np,p,p,p->", G, G, S, S, W)
term2 = jnp.einsum("nk,kp,p,nk,kp->", C, B, W, C, B)
term3 = -2 * jnp.einsum("np,p,p,nk,kp->", G, S, W, C, B)

wgt_ss = term1 + term2 + term3

Computing term1 throws:

1. NotImplementedError: bcoo_spdot_general: batch_dims must correspond to batch dimensions of the sparse representation., if using sparse.BCOO or,
2. NotImplementedError: sparse rule for reduce_sum is not implemented, if using sparse.BCSR

I -could- replace this with jnp.einsum("np,p,p,p", G**2, S, S, W), but I'd rather not duplicate the entire G matrix in-memory, as (n,p) can be quite large.

So I have two questions:

1. Is the above error due to a bug, or lack of implementation details due to sparse/sparse matrix multiplication, and
2. Does einsum even save memory in this case, as opposed to the G**2 approach?

[jakevdp]

The NotImplemented errors you're seeing are due to functionality that has not yet been implemented; I wouldn't really consider this a bug, it's just code that's not yet been written and the error is explicitly stating that.

Regarding performance: to be quite honest, the jax sparse implementations (and in particular sparse-sparse matmul) is quite slow, so I suspect writing G ** 2 vs G, G is going to be the least of your worries 😁

We are working on ideas about how to make this all better, but nothing is ready for production yet.

[jakevdp]

Actually, I suspect that G ** 2 will be much better performance-wise than G, G: the reason is that in your einsum statement, G, G essentially ends up computing G * G, and element-wise multiplication between two sparse matrices requires a set intersection operation between the specified indices of the matrices. By contrast, G ** 2 is basically a single vectorized operation over the defined data, which should be much faster.

[quattro]

Thanks Jake, I agree that the G ** 2 would be more performant, but since I'm computing a scalar statistic, I'd like to avoid creating a temporary n,p matrix in memory. I initially tried to do this with a lax.scan but also received implementation error on not being able to pull out column-specific sub-matrices. I appreciate your fast response as always, and understand completely that sparse functionality is still very much a work in progress.