jax.value_and_grad() returns jax.interpreters.ad.JVPTracer' for parameter

[Cai-fx]

Hi. I'm trying to write a nn.Module where the parameters are a sparse coo matrix.

def mat_coomat_dot(mat: jnp.array, 
                   coo_data: jnp.array,
                   coo_rows: jnp.array,
                   coo_cols: jnp.array,
                   outdim: int
                   ):
    def vec_coomat_dot(vec, coo_data, coo_rows, coo_cols, outdim):
        updates = jnp.multiply(coo_data, vec[coo_rows])
        return jnp.zeros(outdim).at[coo_cols].add(updates)
    return jax.vmap(vec_coomat_dot, in_axes=(0,None, None, None, None), out_axes=0)(mat, coo_data, coo_rows, coo_cols, outdim)


class sparseWLayer(nn.Module):
    features : int
    coo_rows : jnp.array   
    coo_cols : jnp.array
    dtype = jnp.float32

    @compact        
    def __call__(self, x):
        weights = self.param('sparse_weights', nn.initializers.ones, len(self.coo_rows))       
        # in jax.value_and_grad()(), returns <class 'jax.interpreters.ad.JVPTracer'>
        x = mat_coomat_dot(x, weights, self.coo_rows, self.coo_cols, self.features)
        return x
    

When running jax.value_and_grad()(params) I get the following error: jax._src.traceback_util.UnfilteredStackTrace: IndexError: Too many indices for array: 1 non-None/Ellipsis indices for dim 0.

Upon examination the parameter sparse_weights is <class 'jax.interpreters.ad.JVPTracer'> instead of DeviceArray. Can someone help me with this? Many thanks!

[jakevdp]

Regarding the variable being a JVPTracer when you print it within the value_and_grad transform: this is expected behavior. JAX uses tracers as standins for DeviceArray objects when transforming functions with jit, vmap, grad, and other transforms. Take a look at How to think in JAX for some background on this.

Regarding the IndexError: it sounds like you're attempting to index a scalar value within the vmap expression, but it's difficult to tell why because your code snippet doesn't show how you're calling the functions. Here's a simpler example of how that can happen:

import jax
import jax.numpy as jnp

@jax.vmap
def f(x):
  print(f"type(x) = {type(x)}")
  print(f"x.shape = {x.shape}")
  return x[0]

x = jnp.arange(4)
f(x)

type(x) = <class 'jax.interpreters.batching.BatchTracer'>
x.shape = ()
---------------------------------------------------------------------------
UnfilteredStackTrace                      Traceback (most recent call last)
[<ipython-input-4-c3e52b5d9249>](https://localhost:8080/#) in <module>
     10 x = jnp.arange(4)
---> 11 f(x)

...

IndexError: Too many indices for array: 1 non-None/Ellipsis indices for dim 0.

Here x is a 1-dimensional array, and when we pass it to a vmapped function, it is traced with a BatchTracer, and each batch has a shape of (), i.e. a scalar. When we attempt to index into that scalar batch, it raises an error for the same reason that jno.array(0)[0] would raise an error: you cannot index a scalar.

Your situation is probably analogous: the solution would be either to pass a higher-dimensional array to your vmapped function, or perhaps remove the vmap transform if it's not giving you the function semantics that you want.

Hope that helps – if you need more specific recommendations for fixing your problem, I'd suggest editing your question to add a minimal reproducible example that others can run to see the exact behavior you're seeing.

Best of luck!

[Cai-fx]

Hi, Jake. Thank you so much for the clear and detailed response. It is indeed the problem of indexing scalar values. I removed thevmap transform and it worked :)