how to compute jvp and vjp for sparse (BCOO) arguments

[stefanozampini]

JAX provides experimental support for grad and value_and_grad for scalar functions with sparse arguments, and jacfwd and jacrev for (possibly) vector functions. Can you provide an example of how to compute jvp and vjp for vector functions with sparse arguments?

[jakevdp]

jacfwd and jacrev are higher-level transformations that support higher-level operands like BCOO arrays and other PyTree objects as inputs.

By contrast, jvp and vjp are lower-level transformations that only accept raw arrays as inputs, and not pytree objects.

Given this, if you want to use jvp or vjp on BCOO matrices, you'll have to do a lot of boilerplate to pull-out the raw arrays. Here's an example. We'll start with a dense jvp:

import jax
import jax.numpy as jnp
import numpy as np

def f(x):
  return x.sum()

x = jnp.arange(6.0).reshape(2, 3)
dx = jnp.ones_like(x)

primals_out, tangents_out = jax.jvp(f, (x,), (dx,))
print(primals_out)
# 15.0
print(tangents_out)
# 6.0

and here's roughly the sparse equivalent:

from jax.experimental import sparse

def f_raw(data, indices, shape=x.shape):
  x_bcoo = sparse.BCOO((data, indices), shape=shape)
  return f(x_bcoo)

x_dot_bcoo = sparse.BCOO((jnp.ones_like(x_bcoo.data), x_bcoo.indices), shape=x_bcoo.shape)

primals = (x_bcoo.data, x_bcoo.indices)
tangents = (x_dot_bcoo.data, np.zeros_like(x_dot_bcoo.indices, dtype=jax.dtypes.float0))
primals_out, tangents_out = jax.jvp(f_raw, primals, tangents)

print(primals_out)
# 15.0
print(tangents_out)
# 5.0

All that boilerplate is effectively what sparse.grad and sparse.jacobian do for you.

[SNMS95]

Hey @jakevdp ,

Does this mean that everytime we have to call f, we have to create the sparse matrix x_bcoo ?
Is there a way around this? Since inside the function f_raw, you have two copies of the same data, this can be quite difficult for large matrices, right?

[jakevdp]

Does this mean that everytime we have to call f, we have to create the sparse matrix x_bcoo?

If you want to use jax.jvp or jax.vjp, then the answer is yes. But these are low-level functions that you should not need to use directly – I'd recommend using higher-level gradient transforms like jax.grad/jax.jacrev, or if you're differentiating with respect to sparse arguments, sparse.grad/sparse.jacrev, in which case this boilerplate is handled for you.