code: def test_boxed_param(self):
      def f(scope, xs):
          def g(scope, x):
              kernel_init = meta.with_partitioning(
                  nn.initializers.ones_init(), ("in", "out")
              )
              kernel = scope.param(
                  "kernel", kernel_init, (x.shape[-1], 2)
              )
              kernel_box = scope.get_variable("params", "kernel")
              self.assertIsInstance(kernel_box, meta.Partitioned)
              self.assertEqual(kernel_box.names, ("in", "out"))
              return x @ kernel
  
          lift.vmap(
              g,
              in_axes=0,
              variable_axes={"params": 0},
              split_rngs={"params": True},
              metadata_params={meta.PARTITION_NAME: "batch"},
          )(scope, xs)
  
      _, variables = init(f)(random.key(0), jnp.zeros((8, 3)))
      self.assertEqual(
          variables["params"]["kernel"].names, ("batch", "in", "out")
      )
  
file_path: flax/tests/core/core_meta_test.py
imports: import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random, sharding
  from jax.experimental import mesh_utils
  from flax import errors
  from flax.core import init, lift, meta, nn
label_desc: Use image node_id 1 for calling the MetaTest obj's underlying member method code with example usage: obj.test_boxed_param() and returns: unknown