code: def test_param_pytree_with_axes(self):
      def init_fn(k, s, d):
          del k
          return {
              "a": jnp.zeros(s, d),
              "b": (jnp.zeros(s, d), jnp.zeros(s, d)),
          }
  
      axes = {
          "a": ("foo", "bar"),
          "b": (("foo", "bar"), ("bar", "foo")),
      }
  
      class ParamTest(nn.Module):
          @nn.compact
          def __call__(self, x):
              foo = partitioning.param_with_axes(
                  "foo", init_fn, (2, 2), x.dtype, axes=axes
              )
              return x + foo["a"]
  
      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=axes),
      )
      logical_axis_names = partitioning.get_axis_names(
          variables["params_axes"]
      )
      expected = freeze(
          {
              "foo": {
                  "a": jax.sharding.PartitionSpec("foo", "bar"),
                  "b": (
                      jax.sharding.PartitionSpec("foo", "bar"),
                      jax.sharding.PartitionSpec("bar", "foo"),
                  ),
              }
          }
      )
      self.assertEqual(logical_axis_names, expected)
  
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 10 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_param_pytree_with_axes() and returns: dict, unknown