code: def test_variable_with_axes_fallback(self, wsc_fn):
      class VarTest(nn.Module):
          @nn.compact
          def __call__(self, x):
              foo = partitioning.variable_with_axes(
                  "test",
                  "foo",
                  jnp.zeros,
                  (2, 2),
                  x.dtype,
                  axes=("foo", "bar"),
                  fallback=partitioning.RulesFallback.NO_CONSTRAINT,
              )
              return x + foo.value
  
      p_rules = (
          # No rule for 'foo':
          ("bar", "data"),
          ("baz", None),
      )
      k = random.key(0)
      x = jnp.ones((2, 2))
      with partitioning.axis_rules(p_rules):
          variables = VarTest().init(k, x)
  
      wsc_fn.assert_not_called()
      self.assertIn("test", variables)
      self.assertIn("test_axes", variables)
      self.assertEqual(
          variables["test_axes"]["foo_axes"],
          partitioning.AxisMetadata(names=("foo", "bar")),
      )
      logical_axis_names = partitioning.get_axis_names(
          variables["test_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 14 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_variable_with_axes_fallback(wsc_fn) and returns: unknown