[sparse matrix] how to jaxify the sparse matrix

[JiahaoYao]

import scipy
import scipy.sparse
from scipy.sparse import random as sp_rand
from tqdm import tqdm, trange
import jax 
from jax.experimental import sparse

my_sp_mat_list = [ sp_rand(2000, 2000, density=0.01) for i in trange(500)] 

my_sp_mat_list_jax = [sparse.BCOO.fromdense( H.toarray() ) for H in tqdm(my_sp_mat_list)]

I wonder if there is a way to directly convert the sparse matrix in numpy to jax.DeviceArray? Because the current way of first converting the numpy sparse matrix to dense and then convert to the sparse matrix seems costly.

[jakevdp]

You can use BCOO.from_scipy_sparse:

my_sp_mat_list_jax = [sparse.BCOO.from_scipy_sparse(H) for H in my_sp_mat_list]

[jakevdp]

Sorry for the delay – I take my weekends seriously 😀 .

Just as vmap will not work on lists of device arrays, vmap will not work on lists of sparse arrays. What you'll need is to concatenate those. Unfortunately we've not yet implemented the sparsification rule for lax.concatenate_p, but it should be doable. Let me see if I can cook up a short example.

[jakevdp]

Here's a very simple sparse matrix concatenation routine that works for your example. I tried to use asserts to make clear which cases would require more logic to implement correctly:

from jax import lax

def concat(arr_list, axis=0):
  assert len(arr_list) > 0, "Cannot concatenate zero arrays."
  assert axis == 0, "only axis=0 is implemented"
  assert len(set(tuple(a.shape) for a in arr_list)) == 1, "all arrays must have the same shape"
  assert len(set(a.nse for a in arr_list)) == 1, "All entries must have the same nse"
  assert all(a.n_sparse == a.ndim for a in arr_list), "Only sparse dimensions supported"

  data = lax.concatenate([a.data[None] for a in arr_list], dimension=0)
  indices = lax.concatenate([a.indices[None] for a in arr_list], dimension=0)

  return sparse.BCOO((data, indices), shape=(len(arr_list), *arr_list[0].shape))

stacked = concat(my_sp_mat_list_jax)

result = vmap(matvec, in_axes=(0, None))(stacked, vec)

Perhaps we can add a sparse rule for lax.concatenate that does something along these lines.

[JiahaoYao]

This is AWESOME, thanks @jakevdp for the prompt reply!

[jakevdp]

#10466 adds a sparsification rule for concatenation, and #10454 makes added dimensions default to batch. Once these two changes land, you should be able to do something like this:

my_sp_mat_list_jax = [sparse.BCOO.from_scipy_sparse( H ) for H in tqdm(my_sp_mat_list)]
stacked = sparse.sparsify(jnp.stack)(my_sp_mat_list_jax)

[JiahaoYao]

Awesome! as always, thanks for being so supportive! Thanks!