code: def test_boxed_variable(self):
      class Bar(nn.Module):
          @nn.compact
          def __call__(mdl_self, x):  # pylint: disable=no-self-argument
              kernel_init = nn.with_partitioning(
                  nn.initializers.ones_init(), ("in", "out")
              )
              kernel = mdl_self.variable(
                  "params",
                  "kernel",
                  kernel_init,
                  mdl_self.make_rng("params"),
                  (x.shape[-1], 2),
              )
              kernel.value += 1.0
              self.assertEqual(
                  kernel.value.sum(), kernel.value.size * 2
              )
              kernel_box = mdl_self.get_variable("params", "kernel")
              self.assertIsInstance(kernel_box, nn.Partitioned)
              self.assertEqual(kernel_box.names, ("in", "out"))
              return x @ kernel.value
  
      class Foo(nn.Module):
          @nn.compact
          def __call__(self, xs):
              return nn.vmap(
                  Bar,
                  in_axes=0,
                  variable_axes={"params": 0},
                  split_rngs={"params": True},
                  metadata_params={nn.PARTITION_NAME: "batch"},
              )(name="bar")(xs)
  
      m = Foo()
      variables = m.init(random.key(0), jnp.zeros((8, 3)))
      self.assertEqual(
          variables["params"]["bar"]["kernel"].names,
          ("batch", "in", "out"),
      )
  
file_path: flax/tests/linen/linen_meta_test.py
imports: import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from jax.experimental import mesh_utils
  from jax.sharding import Mesh, PartitionSpec
  from flax import linen
label_desc: Use image node_id 2 for calling the LinenMetaTest obj's underlying member method code with example usage: obj.test_boxed_variable() and returns: unknown, nn