code: def test_logical_to_mesh_axes(self):
      axes_0 = ("foo", "bar")
      # direct rule assignment
      self.assertEqual(
          partitioning.logical_to_mesh_axes(axes_0, rules=AXIS_RULES_1),
          ("data", "model"),
      )
      # axis rules context
      with partitioning.axis_rules(AXIS_RULES_1):
          self.assertEqual(
              partitioning.logical_to_mesh_axes(axes_0),
              ("data", "model"),
          )
          # nested context
          with partitioning.axis_rules(AXIS_RULES_2):
              self.assertEqual(
                  partitioning.logical_to_mesh_axes(axes_0),
                  ("model", None),
              )
      # duplicated logical names
      with partitioning.axis_rules(AXIS_RULES_1):
          with self.assertRaises(ValueError):
              partitioning.logical_to_mesh_axes(("foo", "foo", "baz"))
  
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 3 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_logical_to_mesh_axes() without return types