Sparse tensordot with sparse outputs?

[DanPuzzuoli]

Hi,

I've just started playing around with the BCOO sparse arrays and I'm having trouble performing the following operation. Basically, I want to compute a linear combination of a list of a sparse matrices, with the list of sparse matrices being specified as a 3d BCOO array, the coefficients a dense 1d array, and I want the output to be another BCOO. To be exact, I want to do:

• Inputs:
  • Linear combination coefficients c = [c1, ..., ck] - given as a 1d dense array
  • List of sparse matrices A = [A1, ..., Ak] - given as a 3d BCOO
• Output:
  • Sparse array containing sum_i c_i A_i

If everything is dense I would just do: jnp.tensordot(c, A, axes=1). I've tried sparsify(jnp.tensordot)(c, A, axes=1), and the output is mathematically correct but it's dense, and I want it to remain sparse.

An alternative approach I've tried to take is using sparsify(jnp.sum). Doing sparsify(jnp.sum)(A, axis=0) does result in a sparse version of sum_i A_i, so I was also hoping for some way to multiply the i^th entry of A with c_i. E.g. if everything is dense, c[:, jnp.newaxis, jnp.newaxis] * A does the trick, but this doesn't work for A being BCOO.

Any ideas? Thanks!

[DanPuzzuoli]

I just realized I can actually do:

sparsify(jnp.sum)(jnp.broadcast_to(c[:, jnp.newaxis, jnp.newaxis], A.shape) * A, axis=0)

though this seems a little excessive as the broadcast_to call creates an enormous dense array.

[jakevdp]

It was added relatively recently, but if you do the (generalized) dot product of two sparse matrices the result will be sparse. Using the main branch, you could do something like this:

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

c = jnp.array(np.random.rand(10))
M = jnp.array(np.random.rand(10, 3, 4))
M = M.at[M < 0.5].set(0)

Msp = sparse.BCOO.fromdense(M)
csp = sparse.BCOO.fromdense(c)

def f(M, c):
  return jnp.tensordot(M, c, axes=(0, 0))

print(f(M, c))
# [[1.0761551  1.1919075  2.4007194  1.0579658 ]
#  [2.5757928  0.77669764 1.0757202  1.4997834 ]
#  [0.14920855 0.8391593  1.8824157  1.1219863 ]]
print(sparse.sparsify(f)(Msp, csp).todense())
# [[1.0761552  1.1919074  2.4007192  1.0579658 ]
#  [2.5757926  0.77669764 1.0757202  1.4997835 ]
#  [0.14920855 0.8391593  1.8824158  1.1219863 ]]

Currently it's limited to cases where one of the matrices has a single sparse dimension, and only sparse dimensions, not batch dimensions, can be contracted.

[DanPuzzuoli]

Any thoughts on if c is dense? Is converting it to sparse via fromdense something to worry about performance wise? Assuming you can jit/grad through this it isn't too big a deal for my specific applications, as c and the first dimension of A won't be that long (but the last two dimensions of A will be extremely large).

[jakevdp]

Another thing you can do here is use element-wise broadcasted multiplication followed by a sum:

result = (c[:, None, None] * Msp).sum(0).todense()

again, I think that requires using the main branch because the broadcasted multiplication has not yet made it into a release.

[jakevdp]

In general, though, I think we'll stick with the convention that sparse-dense dot products output dense, and that sparse-sparse dot products output sparse. When sparse-sparse matmul is fully implemented, you should be able to create a sparse matrix with all dense dimensions (i.e. something like BCOO.fromdense(x, n_dense=x.ndim), then use this to do sparse-dense->sparse matmuls. But unfortunately contraction over dense dimensions in spdot_general is not yet implemented.

It's on my list though 😁

[DanPuzzuoli]

Thanks for the detailed replies.

As I'm working on a feature in a package that uses JAX as a dependency (and so want to stick to things already in a release), I think for now I'll proceed with:

sparsify(jnp.sum)(jnp.broadcast_to(c[:, None, None], A.shape) * A, axis=0)

as I'm already seeing a substantial speed improvement over dense arrays (on CPU) in a toy example. I'll make sure to update it once the element-wise broadcasted multiplication makes it into a release.

Thanks for the help Jake! JAX is awesome!