Using pmap with a sparse matrix

[IrishWhiskey]

Hey,

I'm trying to use a batched COO sparse matrix with pmap but I get some errors. For example, take a look at the code below:

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

def f(x, y):
    return jnp.dot(x, y)

x = jnp.array(np.random.rand(8, 1000))
y = jnp.array(np.random.rand(1000))

x_bcoo = BCOO.fromdense(x)

jax.experimental.sparse.sparsify(jax.pmap(f, in_axes=(0, None)))(x_bcoo, y)

This fails with the following error message:

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
/tmp/ipykernel_15251/3140807170.py in <cell line: 3>()
      1 x_bcoo = BCOO.fromdense(x)
      2 
----> 3 jax.experimental.sparse.sparsify(jax.pmap(f, in_axes=(0, None)))(x_bcoo, y)

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/transform.py in wrapped(*args, **params)
    426     spenv = SparsifyEnv()
    427     spvalues = arrays_to_spvalues(spenv, args)
--> 428     spvalues_out, out_tree = f_raw(spenv, *spvalues, **params)
    429     out = spvalues_to_arrays(spenv, spvalues_out)
    430     return tree_unflatten(out_tree, out)

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/transform.py in wrapped(spenv, *spvalues, **params)
    412     wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(f, params), in_tree)
    413     jaxpr, out_avals_flat, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals_flat)
--> 414     result = eval_sparse(jaxpr, consts, spvalues_flat, spenv)
    415     if len(out_avals_flat) != len(result):
    416       raise Exception("Internal: eval_sparse does not return expected number of arguments. "

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/transform.py in eval_sparse(jaxpr, consts, spvalues, spenv)
    384     if any(val.is_sparse() for val in invals):
    385       if prim not in sparse_rules:
--> 386         _raise_unimplemented_primitive(prim)
    387       out = sparse_rules[prim](spenv, *invals, **eqn.params)
    388     else:

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/transform.py in _raise_unimplemented_primitive(primitive)
    118                               "would result in dense output. If this is your intent, use "
    119                               "sparse.todense() to convert your arguments to dense matrices.")
--> 120   raise NotImplementedError(f"sparse rule for {primitive} is not implemented.")
    121 
    122 

NotImplementedError: sparse rule for xla_pmap is not implemented.

Instead, If I run the following:

def f_sparse(x, y):
    return jnp.dot(x.todense(), y)

jax.pmap(f_sparse, in_axes=(0, None))(x_bcoo, y)

I get this other error message:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_15251/2767757453.py in <cell line: 4>()
      2     return jnp.dot(x.todense(), y)
      3 
----> 4 jax.pmap(f_sparse, in_axes=(0, None))(x_bcoo, y)

    [... skipping hidden 12 frame]

/tmp/ipykernel_15251/2767757453.py in f_sparse(x, y)
      1 def f_sparse(x, y):
----> 2     return jnp.dot(x.todense(), y)
      3 
      4 jax.pmap(f_sparse, in_axes=(0, None))(x_bcoo, y)

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/bcoo.py in todense(self)
   2563   def todense(self) -> Array:
   2564     """Create a dense version of the array."""
-> 2565     return bcoo_todense(self)
   2566 
   2567   def transpose(self, axes: Optional[Sequence[int]] = None) -> BCOO:

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/bcoo.py in bcoo_todense(mat)
    180     mat_dense: dense version of ``mat``.
    181   """
--> 182   return _bcoo_todense(mat.data, mat.indices, spinfo=mat._info)
    183 
    184 def _bcoo_todense(data: Array, indices: Array, *, spinfo: BCOOInfo) -> Array:

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/bcoo.py in _bcoo_todense(data, indices, spinfo)
    195     mat : array with specified shape and dtype matching ``data``
    196   """
--> 197   return bcoo_todense_p.bind(jnp.asarray(data), jnp.asarray(indices), spinfo=spinfo)
    198 
    199 @bcoo_todense_p.def_impl

    [... skipping hidden 5 frame]

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/bcoo.py in _bcoo_todense_abstract_eval(data, indices, spinfo)
    218 def _bcoo_todense_abstract_eval(data, indices, *, spinfo):
    219   shape = spinfo.shape
--> 220   _validate_bcoo(data, indices, shape)
    221   return core.ShapedArray(shape, data.dtype)
    222 

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/bcoo.py in _validate_bcoo(data, indices, shape)
    121 
    122 def _validate_bcoo(data: Buffer, indices: Buffer, shape: Sequence[int]) -> BCOOProperties:
--> 123   props = _validate_bcoo_indices(indices, shape)
    124   n_batch, n_sparse, n_dense, nse = props
    125   shape = tuple(shape)

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/jax/experimental/sparse/bcoo.py in _validate_bcoo_indices(indices, shape)
    134   assert jnp.issubdtype(indices.dtype, jnp.integer)
    135   shape = tuple(shape)
--> 136   nse, n_sparse = indices.shape[-2:]
    137   n_batch = len(indices.shape) - 2
    138   n_dense = len(shape) - n_batch - n_sparse

ValueError: not enough values to unpack (expected 2, got 1)

Am I doing anything wrong? Does pmap support sparse matrices? (If not, is this a feature that will be added in the near future?)

I'm using the following jax and jaxlib versions

jax==0.4.1
jaxlib==0.4.1+cuda11.cudnn86

[jakevdp]

Thanks for the report – this looks like a bug. I'm not aware of anyone attempting to use jax.experimental.sparse with pmap, so I suspect this is simply not possible in the current version of JAX, but I'll take a look to see if I can fix the issue.

[jakevdp]

This is the issue: #13932

[IrishWhiskey]

Thank you Jake.