code: def test_get_sharding(self):
      devices = mesh_utils.create_device_mesh(
          (jax.local_device_count(), 1)
      )
      mesh = sharding.Mesh(devices, ("in", "out"))
      xs = {
          "kernel": meta.Partitioned(jnp.zeros((8, 3)), ("in", "out")),
          "bias": jnp.zeros((8, 3)),
          "step": jnp.array(100),
      }
      ps = meta.get_sharding(xs, mesh)
      self.assertEqual(
          ps,
          {
              "kernel": jax.sharding.NamedSharding(
                  mesh, jax.sharding.PartitionSpec("in", "out")
              ),
              "bias": jax.sharding.NamedSharding(
                  mesh, jax.sharding.PartitionSpec()
              ),
              "step": jax.sharding.NamedSharding(
                  mesh, jax.sharding.PartitionSpec()
              ),
          },
      )
  
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 7 for calling the MetaTest obj's underlying member method code with example usage: obj.test_get_sharding() without return types