code: def test_param_with_axes(self):
      class ParamTest(nn.Module):
          @nn.compact
          def __call__(self, x):
              foo = partitioning.param_with_axes(
                  "foo",
                  lambda k, s, d: jnp.zeros(s, d),
                  (2, 2),
                  x.dtype,
                  axes=("foo", "bar"),
              )
              return x + foo
  
      p_rules = (("foo", "model"), ("bar", "data"), ("baz", None))
      k = random.key(0)
      x = jnp.ones((2, 2))
      with partitioning.axis_rules(p_rules):
          variables = ParamTest().init(k, x)
      self.assertIn("params", variables)
      self.assertIn("params_axes", variables)
      self.assertEqual(
          variables["params_axes"]["foo_axes"],
          partitioning.AxisMetadata(names=("foo", "bar")),
      )
      logical_axis_names = partitioning.get_axis_names(
          variables["params_axes"]
      )
      self.assertEqual(
          logical_axis_names,
          {"foo": jax.sharding.PartitionSpec("foo", "bar")},
      )
  
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 9 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_param_with_axes() and returns: unknown