Jax vmap - ValueError: vmap must have at least one non-None value in in_axes

[GarfieldGa]

@partial(jax.jit, static_argnames=['self'])
    def forward(self, params: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
        """ Prediction function
            run one forward pass on given input data
            --- input ---
            x: input data of shape (nmode, dim)
            --- return ---
            predicted score vector of size (num_class,)
        """
        # do stuff
        return y_hat
    v_forward = jax.vmap(forward, in_axes=(None,None,0)) # returns (N, num_class)

I have the above code defined in class A, and I have a class B inherited from A, i.e., class B(A). The function forward and v_forward are not overwritten in B. Then when I tried to call self.v_forward(arg1, arg2), it gives an error saying
ValueError: vmap must have at least one non-None value in in_axes

Does anyone know how this comes, since I do have a 0 in my in_axies tuple.

[jakevdp]

Hi – I tried to reproduce your issue based on the description, but I can't. Here's the code I wrote:

import jax

class A:
  def f(self, x, y):
    return x + y
  vf = jax.vmap(f, in_axes=(None, None, 0))

class B(A):
  pass

A().vf(1, jax.numpy.arange(4)) # No error
B().vf(1, jax.numpy.arange(4)) # No error

Can you please update your question with a minimal reproducible example that recreates the error you're seeing? Thanks!

[GarfieldGa]

Thanks for the reply, by reproducing the error, it actually let me find the issue. The problem was with the way I pass in the arguments. I had some other codes producing the third argument for vf method in your example. The third argument was 'None' due to some mistake which gives rise to the error above.

[jakevdp]

Great! A side note: I noticed you marking self as static in your example. In general this is a bad idea; for more discussion of this (and how to correctly JIT-compile methods) see JAX FAQ: how to use JIT with methods.