code: def test_basic(self):
      class Block(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs):
              self.linear = nnx.Linear(3, 3, rngs=rngs)
  
          def __call__(self, x: jax.Array) -> jax.Array:
              x = self.linear(x)
              x = nnx.gelu(x)
              return x
  
      MLP = nnx.Vmap(Block, variable_axes={nnx.Param: 0}, axis_size=5)
  
      module = MLP(rngs=nnx.Rngs(0))
  
      assert not jnp.allclose(
          module.vmap_module.linear.kernel[0],
          module.vmap_module.linear.kernel[1],
      )
      assert module.vmap_module.linear.kernel.shape == (5, 3, 3)
      assert module.vmap_module.linear.bias.shape == (5, 3)
  
      x = jnp.ones((5, 1, 3))
      y = module(x)
  
      assert y.shape == (5, 1, 3)
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestVmap obj's underlying member method code with example usage: obj.test_basic() and returns: x