Flax models with sparse matrices

[IrishWhiskey]

I want to use sparse matrices with a Flax model so I tried to transform the model.apply function with sparse.sparsify. Unfortunately when I run the resulting function I get an error. Is there any way to make it work?

Here is the code I ran:

import jax
from jax.experimental import sparse

import flax
import flax.linen as nn

class Model(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.relu(nn.Dense(64)(x))
        x = nn.Dense(1)(x)
        return x

model = Model()
data = jax.random.normal(jax.random.PRNGKey(42), shape=(1000, 64))
params = model.init(jax.random.PRNGKey(0), data)

data_sp = sparse.BCOO.fromdense(data)
sparse.sparsify(model.apply)(params, data_sp)

Here is the error message I get:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[6], line 1
----> 1 sparse.sparsify(model.apply)(params, data_sp)

File ~/anaconda3/envs/tensorflow2_p310/lib/python3.10/site-packages/jax/experimental/sparse/transform.py:466, in _sparsify_with_interpreter.<locals>.wrapped(*args, **params)
    464 spenv = SparsifyEnv()
    465 spvalues = arrays_to_spvalues(spenv, args)
--> 466 spvalues_out, out_tree = f_raw(spenv, *spvalues, **params)
    467 out = spvalues_to_arrays(spenv, spvalues_out)
    468 return tree_unflatten(out_tree, out)

File ~/anaconda3/envs/tensorflow2_p310/lib/python3.10/site-packages/jax/experimental/sparse/transform.py:451, in sparsify_raw.<locals>.wrapped(spenv, *spvalues, **params)
    449 wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(f, params), in_tree)
    450 jaxpr, out_avals_flat, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals_flat)
--> 451 result = eval_sparse(jaxpr, consts, spvalues_flat, spenv)
    452 if len(out_avals_flat) != len(result):
    453   raise Exception("Internal: eval_sparse does not return expected number of arguments. "
    454                   "Got {result} for avals {out_avals_flat}")

File ~/anaconda3/envs/tensorflow2_p310/lib/python3.10/site-packages/jax/experimental/sparse/transform.py:430, in eval_sparse(jaxpr, consts, spvalues, spenv)
    428   out = sparse_rules_bcoo[prim](spenv, *invals, **eqn.params)
    429 else:
--> 430   out_bufs = prim.bind(*(spenv.data(val) for val in invals), **eqn.params)
    431   out_bufs = out_bufs if prim.multiple_results else [out_bufs]
    432   out = []

TypeError: CustomJVPCallPrimitive.bind() got an unexpected keyword argument 'call_jaxpr'

I'm using Flax 0.7.5, JAX 0.4.23 on Python 3.10.

[chiamp]

It seems like the error stems from the relu function. If you try something like sparse.sparsify(nn.relu)(data_sp), you will get an error: NotImplementedError: sparse rule for max is not implemented.
If you remove the relu from your Model, the code will run without issues. I am currently speaking with the Jax team to see if there's a way to use sparsify in conjunction with relu