Instead of casting PartitionSpec Tuples to standard Tuples, convert standard

Tuples to PartitionSpecs for comparision statements.

Author: QuinnMMcGarry

axlearn/common/attention_test.py

@@ -6433,7 +6433,7 @@ def test_logit_sink_parameter_initialization(self):
         self.assertIn("sink", param_specs)
         sink_spec = param_specs["sink"]
         self.assertEqual(sink_spec.shape, (num_heads,))
-        self.assertEqual(tuple(sink_spec.mesh_axes), ("model",))
+        self.assertEqual((sink_spec.mesh_axes), jax.P("model",))
         self.assertEqual(sink_spec.weight_decay_scale, 0.0)
 
     def test_logit_sink_disabled_by_default(self):

axlearn/common/flash_attention/layer_test.py

@@ -526,7 +526,7 @@ def as_partition_spec(pytree: CompositeAttentionBias) -> PartitionSpec:
             spec = test_layer._logit_biases_spec(segment_ids)  # pylint: disable=protected-access
             spec = as_partition_spec(spec)
             self.assertIsInstance(spec, PartitionSpec)
-            self.assertEqual(tuple(spec), test_layer.config.mha_dim_to_partition_spec["btnh"][:2])
+            self.assertEqual(spec, jax.P(test_layer.config.mha_dim_to_partition_spec["btnh"][:2]))
 
     @parameterized.product(
         _TEST_CONFIGS,

axlearn/common/layers_test.py

@@ -572,14 +572,14 @@ def test_rms_norm_partition_specs_constraint(self, mock_with_sharding_constraint
 
         # 1. Input tensor constraint.
         input_spec = calls[0].args[1]
-        self.assertEqual(tuple(input_spec), ("fsdp", "model", None))
+        self.assertEqual((input_spec), jax.P("fsdp", "model", None))
         self.assertEqual(calls[0].args[0].shape, (2, 3, dim))
         self.assertEqual(calls[0].args[0].dtype, jnp.float32)
         np.testing.assert_array_equal(calls[0].args[0], inputs)
 
         # 2. Output tensor constraint.
         output_spec = calls[1].args[1]
-        self.assertEqual(tuple(output_spec), ("fsdp", None, None))
+        self.assertEqual((output_spec), jax.P("fsdp", None, None))
         self.assertEqual(calls[1].args[0].shape, (2, 3, dim))
         self.assertEqual(calls[1].args[0].dtype, jnp.float32)
         np.testing.assert_array_equal(calls[1].args[0], outputs)

axlearn/common/lora_test.py

@@ -50,8 +50,8 @@ def test_set_param_spec_config(self):
         layer_cfg.lora_up.param_partition_spec = [None, "data"]
         layer = layer_cfg.instantiate(parent=None)
         param_specs = layer.create_parameter_specs_recursively()
-        self.assertEqual(tuple(param_specs["lora_down"]["weight"].mesh_axes), ("data", None))
-        self.assertEqual(tuple(param_specs["lora_up"]["weight"].mesh_axes), (None, "data"))
+        self.assertEqual((param_specs["lora_down"]["weight"].mesh_axes), jax.P("data", None))
+        self.assertEqual((param_specs["lora_up"]["weight"].mesh_axes), jax.P(None, "data"))
 
     def test_forward(self):
         input_dim = 2