Is nnx.vmap a replacement for jax.vmap? Getting different behavior

[jlperla]

Take the following simple example:

import jax
import jax.numpy as jnp
from jax import random
from flax import nnx
N = 10
M = 2
sigma = 0.001
rngs = nnx.Rngs(42)
theta = random.normal(rngs(), (M,))
X = random.normal(rngs(), (N, M))
Y = X @ theta + sigma * random.normal(rngs(), (N,)) 

model = nnx.Linear(M, 1, use_bias=False, rngs=rngs)

def residual(model, x, y):
    y_hat = model(x)
    return (y_hat - y) ** 2

residual(model, X[0], Y[0])
jax.vmap(residual, in_axes=(None, 0, 0))(model, X, Y)

# Fails? Shouldn't just be a replacement?
# nnx.vmap(residual, in_axes=(None, 0, 0))(model, X, Y)

Note the last line with nnx.vmap which returns

  File "/Users/jlperla/anaconda3/envs/econ622/lib/python3.11/site-packages/flax/nnx/nnx/transforms/parallelization.py", line 305, in vmap_wrapper
    raise ValueError(
ValueError: Inconsistent lengths between state_axes states and arguments: {10, 2}

Am I misunderstanding this function? I am not entirely sure I need to use it in the circumstances, but I am having trouble forming a mental model of why these are different calls?

[cgarciae]

@jlperla your example now works in flax>=0.9.0 which introduced JAX-style transforms (see #4107). In previous versions the vectorization configuration for all Modules was controlled by a state_axes argument. To fix the example in your current version you can do this:

jax.vmap(residual, in_axes=(None, 0, 0), state_axes={...: None})(model, X, Y)

But I'd suggest upgrading to the new API instead and keep your example as is.

[cgarciae]

Am I misunderstanding this function? I am not entirely sure I need to use it in the circumstances, but I am having trouble forming a mental model of why these are different calls?

I'm very happy we made the change to JAX-style transforms, starting from 0.9.0 the mental model for Flax transforms is the same as for JAX transforms which should help users a lot.

[jlperla]

Verified that this works on 0.9 on both windows and macos!