code: def test_boxed_param_with_mesh(self):
      devices = mesh_utils.create_device_mesh(
          (jax.local_device_count(), 1)
      )
      mesh = sharding.Mesh(devices, ("in", "out"))
  
      def f(scope, x):
          kernel_init = meta.with_partitioning(
              nn.initializers.ones_init(), ("in", "out"), mesh=mesh
          )
          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
  
      @jax.jit
      def create_state():
          y, variables = init(f)(random.key(0), jnp.zeros((8, 4)))
          spec = meta.get_partition_spec(variables)
          shardings = jax.tree_map(
              lambda s: sharding.NamedSharding(mesh, s), spec
          )
          variables = jax.lax.with_sharding_constraint(
              variables, shardings
          )
          return variables
  
      variables = create_state()
      self.assertEqual(
          variables["params"]["kernel"].names, ("in", "out")
      )
      self.assertIs(variables["params"]["kernel"].mesh, mesh)
  
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 8 for calling the MetaTest obj's underlying member method code with example usage: obj.test_boxed_param_with_mesh() and returns: unknown, variables