code: def test_axis_rules(self):
      self.assertEqual(partitioning._axis_rules.rules, ())
      partitioning.set_axis_rules(AXIS_RULES_1)
      self.assertEqual(partitioning._axis_rules.rules, AXIS_RULES_1)
      self.assertEqual(partitioning.get_axis_rules(), AXIS_RULES_1)
      partitioning.set_axis_rules(())
  
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 1 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_axis_rules() without return types