code: def test_vmap_with_axes(self):
      class Foo(nn.Module):
          @nn.compact
          def __call__(self, x):
              return (
                  partitioning.param_with_axes(
                      "w",
                      jax.nn.initializers.uniform(),
                      [4, 3],
                      axes=("out", "in"),
                  )
                  @ x
              )
  
      class Vmapped(nn.Module):
          @nn.compact
          def __call__(self, x):
              FooVmapped = partitioning.vmap_with_axes(  # pylint: disable=invalid-name
                  Foo,
                  variable_axes={
                      "params": 1,
                  },
                  split_rngs={"params": True},
                  partitioning_axis_names={"params": "vmap_axis"},
              )
              return FooVmapped(name="foo_vmapped")(x)
  
      p_rules = (("out", None), ("in", "data"), ("vmap_axis", "model"))
  
      # check that regular Food module is correct
      with partitioning.axis_rules(p_rules):
          variables = Foo().init(
              jax.random.key(0), jnp.array([1, 2, 3])
          )
      variables = unfreeze(variables)
      variables["params"] = jax.tree_util.tree_map(
          lambda x: x.shape, variables["params"]
      )
      self.assertDictEqual(
          variables,
          {
              "params": {"w": (4, 3)},
              "params_axes": {
                  "w_axes": partitioning.AxisMetadata(
                      names=("out", "in")
                  )
              },
          },
      )
  
      # check that FooVmapped adds 'vmap_axis' to axis 1
      with partitioning.axis_rules(p_rules):
          variables = Vmapped().init(
              jax.random.key(0), jnp.array([[1, 2, 3], [4, 5, 6]])
          )
      variables = unfreeze(variables)
      variables["params"] = jax.tree_util.tree_map(
          lambda x: x.shape, variables["params"]
      )
      self.assertDictEqual(
          variables,
          {
              "params": {"foo_vmapped": {"w": (4, 2, 3)}},
              "params_axes": {
                  "foo_vmapped": {
                      "w_axes": partitioning.AxisMetadata(
                          names=("out", "vmap_axis", "in")
                      )
                  }
              },
          },
      )
  
imports: import jax
  import jax.numpy
  from absl.testing import absltest, parameterized
  from jax import random, sharding
  from jax.experimental import mesh_utils
  from flax import linen
  from flax.core import freeze, unfreeze
  from flax.linen import partitioning
label_desc: Use image node_id 16 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_vmap_with_axes() and returns: unknown, FooVmapped