Use nnx vmap with graph True

Author: samanklesaria

tests/nnx/prefix_test.py

@@ -14,11 +14,11 @@ def __init__(self, rngs: nnx.Rngs):
 
       rngs = nnx.Rngs(0, dropout=jax.random.key(1))
       rngs = rngs.split({f'dropout': 5})
-      model = jax.vmap(Model, in_axes=(nnx.prefix(rngs, {'dropout': 0}),))(rngs)
+      model = nnx.vmap(Model, graph=False, graph_updates=False, in_axes=(nnx.prefix(rngs, {'dropout': 0}),))(rngs)
       assert model.drop.rngs.key[...].shape == (5,)
       assert model.drop.rngs.count[...].shape == (5,)
       bias = model.linear.bias[...]
       assert all(jnp.allclose(x,y) for (x,y) in zip(bias, bias[1:]))
 
     # Problem: need nnx.vmap to work with prefix
-    # Currently, vmap flattens to a graphdef first. We must disable this.
+    # The out_axes argument in nnx.vmap is breaking things.