How to implement pytorch's sampled_addmm in jax?

[kkew3]

I want to do very large dense-dense matrix multiplication but only on very few elements of the result matrix. So the overall operation is (dense x dense) * sparsity -> sparse. In pytorch, this can be implemented via torch.sparse.sampled_addmm. How to implement this efficiently in jax? Thank you so much!

[jakevdp]

JAX doesn't have any fully supported sparse computation, but there is jax.experimental.sparse, and bcoo_dot_general_sampled is the equivalent operation you're looking for.

That said, you shouldn't expect that function to be particularly performant: it's really just a reference implementation for exploring the interaction between sparse computation and JAX transformations.

With that caveat, here's what it would look like to call that API:

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

x = jnp.arange(6).reshape(2, 3)
y = jnp.arange(12).reshape(3, 4)

mask = jnp.array([[1, 0, 0, 1],
                  [0, 1, 0, 1]])
result_direct = (x @ y) * mask
print(result_direct)
# [[20  0  0 29]
#  [ 0 68  0 92]]

indices = sparse.BCOO.fromdense(mask).indices
data = sparse.bcoo_dot_general_sampled(
    x, y, indices, dimension_numbers=(([1], [0]), ([], [])))
result_sparse = sparse.BCOO((data, indices), shape=mask.shape)
print(result_sparse.todense())
# [[20  0  0 29]
#  [ 0 68  0 92]]