Add prefix test when calling

Author: samanklesaria

tests/nnx/rngs_test.py

@@ -221,6 +221,8 @@ class Model(nnx.Module):
       def __init__(self, rngs: nnx.Rngs):
         self.linear = nnx.Linear(20, 10, rngs=rngs)
         self.drop = nnx.Dropout(0.1, rngs=rngs)
+      def __call__(self, x):
+        return self.drop(self.linear(x))
 
     with nnx.graphlib.set_graph_updates(False):
       with nnx.graphlib.set_graph_mode(False):
@@ -233,6 +235,10 @@ def __init__(self, rngs: nnx.Rngs):
         bias = model.linear.bias[...]
         assert all(jnp.allclose(x,y) for (x,y) in zip(bias, bias[1:]))
 
+        # This is the same as just using 0 as the model in_axes
+        prefix2 = nnx.prefix(model, {nnx.Variable: 0})
+        nnx.vmap(Model.__call__, in_axes=(prefix2,None))(model, jnp.ones(20))
+
 
   def test_random_helpers(self):
     rngs = nnx.Rngs(0, params=1)