Sparse-Sparse Matrix Multiplication Gives Incorrect nse for n_batch=0

[CamRuiz]

I am new to JAX and was trying out the experimental.sparse module to see if doing sparse-sparse matrix multiplication was feasible on the GPU. However, I ran into an issue where the nse value after multiplying the matrices was much larger than it should have been. For example, when doing A @ B where both are BCOO matrices and A is diagonal, the resulting nse was A.nse * B.nse when in theory it only needs to be equal to B.nse. I found that it was possible to get the correct behaviour when using n_batch=1 for A. Here is some code replicating the issue:

import numpy as np
import scipy as sc

import jax
from jax import numpy as jnp
from jax.experimental import sparse as jsparse
import jax.config as config; config.update("jax_enable_x64", True)

# jax.config.update('jax_platform_name', 'cpu')

np.random.seed(0)

R = sc.sparse.csc_matrix(sc.stats.bernoulli.rvs(size=(1000, 16), p=0.001))

Q_vec = np.random.normal(size=1000)
Q = sc.sparse.dia_matrix((Q_vec, [0]), shape=(1000, 1000))


Q@R
# <1000x16 sparse matrix of type '<class 'numpy.float64'>'
# 	with 16 stored elements in Compressed Sparse Row format>

sparse_matmul = jsparse.sparsify(jnp.matmul)

R_jax = jsparse.BCOO.from_scipy_sparse(R)
Q_jax = jsparse.BCOO.fromdense(Q.todense())
sparse_matmul(Q_jax, R_jax)
# BCOO(float64[1000, 16], nse=16000)

Q_jax = jsparse.BCOO.fromdense(Q.todense(), n_batch=1)
# BCOO(float64[1000, 16], nse=16, n_batch=1)

From what I have read on various other issues it seems that it is necessary to set n_batch to be non-zero when doing sparse-sparse multiplication but would like to check if this is correct? I am also curious if there are any plans to change this so that it works without batching, or if this is possible with the BCOO format or other formats that may be added to JAX in the future.

Furthermore, I've seen that it is not possible to set n_batch when using BCOO.from_scipy_sparse. This would be quite useful for me so I was wondering when this may be added?

[jakevdp]

Thanks for the question: the expected nse in this case is A.nse * B.nse, because that's the upper bound of the required nse for arbitrary sparse matrices. To see why, consider a situation like this:

import jax.numpy as jnp

# A is 4x4 with nse=3
A = jnp.array([[1, 0, 0, 0],
               [1, 0, 0, 0],
               [1, 0, 0, 0],
               [0, 0, 0, 0]])

# B is 4 x 5 with nse=4
B = jnp.array([[1, 1, 1, 1, 0],
               [0, 0, 0, 0, 0],
               [0, 0, 0, 0, 0],
               [0, 0, 0, 0, 0]])

# C is 4 x 5 with nse = A.nse * B.nse = 12
C = A @ B
print(C)
# [[1 1 1 1 0]
#  [1 1 1 1 0]
#  [1 1 1 1 0]
#  [0 0 0 0 0]]

In the worst case, C.nse will be equal to A.nse * B.nse. JAX's JIT computation model requires statically sized outputs (i.e. the shape of output arrays cannot depend on the contents of input arrays). For sparse matrices, this means that the shape of output arrays cannot depend on the indices of the stored values, only on the number of stored values. So in the case of matrix multiplication, the output must be able to handle this worst case, therefore C.nse must be A.nse * B.nse, and in many cases this will be an overestimate of the strict number required. Does that make sense?

[CamRuiz]

Hi Jake,

Thanks for responding so quickly and for your example, that makes perfect sense.