(Even immutable) vmap transformed module cannot be passed by closure?

[y1xiaoc]

Hello, I recently update the flax version and got an error of the following
ValueError: Cannot extract graph node from different trace level, got ...
It seems that I can isolate the problem to passing by closure an nnx module that contains nnx.vmap (either for its methods or from its submodule). For example, take the example in the guide here:

class WeightStack(nnx.Module):
  @nnx.vmap
  def __init__(self, seed: jax.Array):
    self.kernel = nnx.Param(jax.random.uniform(jax.random.key(seed), (2, 3)))
    self.bias = nnx.Param(jnp.zeros((3,)))

  @nnx.vmap(in_axes=0, out_axes=1)
  def __call__(self, x: jax.Array):
    assert self.kernel.ndim == 2, 'Batch dimensions not allowed'
    assert x.ndim == 1, 'Batch dimensions not allowed'
    return x @ self.kernel + self.bias

weights = WeightStack(jnp.arange(10))

x = jax.random.normal(jax.random.key(1), (10, 2))
y = weights(x)

this code above works just fine, however, if one tries to jit the weight, or use weight in any other function and jit that function, the above error would raise.

# value error
nnx.jit(weights)(x)

or

# value error
def build_apply(model):
    return lambda x: model(x)
apply_fn = build_apply(weights)
nnx.jit(apply_fn)(x)

would both throw error like

ValueError: Cannot extract graph node from different trace level, got Param( # 30 (240 B)
  value=Array([[0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.]], dtype=float64)
)

This is quite inconvenient because that would require any transform (that could be in the middle of a large function) to pass the module explicitly, even though the module is not mutating anything, leading to a significant refactoring. I noticed that this error only happens after 0.10.5 update and for now I'm rolling back to 0.10.4. This could also be related to a previous discussion #4804

I'm wondering if this means that nnx.vmap internally mutate the state of the module. If so, can we have a separate "static" vmap that does not have this issue? For those who only requires a batched dimension such static version should be sufficient, and would not have the problem of unable to be passed by closure.

[y1xiaoc]

In fact, instead of putting nnx.vmap on the __call__ function, writing a new custom vmap like the following would do the work

def vmap_module(module, *vmap_args, **vmap_kwargs):
    graph, state = nnx.split(module)
    def inner_apply(state, *args, **kwargs):
        module = nnx.merge(graph, state)
        return module(*args, **kwargs)
    vmap_apply = jax.vmap(inner_apply, *vmap_args, **vmap_kwargs)
    def apply(*args, **kwargs):
        return vmap_apply(state, *args, **kwargs)
    return apply

y = jax.jit(vmap_module(weights))(x) # or nnx.jit, both would work