Casting PartitionSepc to Tuple for b/505613482

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(sink_spec.mesh_axes, ("model",))
+        self.assertEqual(tuple(sink_spec.mesh_axes), ("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(spec, test_layer.config.mha_dim_to_partition_spec["btnh"][:2])
+            self.assertEqual(tuple(spec), 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(input_spec, ("fsdp", "model", None))
+        self.assertEqual(tuple(input_spec), ("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(output_spec, ("fsdp", None, None))
+        self.assertEqual(tuple(output_spec), ("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)
@@ -1462,21 +1462,21 @@ def test_embed_partition_specs_constraint(self, mock_with_sharding_constraint):
 
         # 1. Input activation constraint (indices tensor).
         input_spec = calls[0].args[1]
-        self.assertEqual(input_spec, ("fsdp", None))
+        self.assertEqual(tuple(input_spec), ("fsdp", None))
         self.assertEqual(calls[0].args[0].shape, (3, seq_len))
         self.assertEqual(calls[0].args[0].dtype, jnp.int32)
         np.testing.assert_array_equal(calls[0].args[0], ixs)
 
         # 2. Embedding weight constraint.
         weight_spec = calls[1].args[1]
-        self.assertEqual(weight_spec, ("model", "fsdp"))
+        self.assertEqual(tuple(weight_spec), ("model", "fsdp"))
         self.assertEqual(calls[1].args[0].shape, (num_embeddings, dim))
         self.assertEqual(calls[1].args[0].dtype, jnp.float32)
         np.testing.assert_array_equal(calls[1].args[0], state["weight"])
 
         # 3. Output activation constraint (after lookup).
         output_spec = calls[2].args[1]
-        self.assertEqual(output_spec, ("fsdp", "model"))
+        self.assertEqual(tuple(output_spec), ("fsdp", "model"))
         self.assertEqual(calls[2].args[0].shape, (3, seq_len, dim))
         self.assertEqual(calls[2].args[0].dtype, jnp.float32)
         np.testing.assert_array_equal(calls[2].args[0], actual_embeds)

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(param_specs["lora_down"]["weight"].mesh_axes, ("data", None))
-        self.assertEqual(param_specs["lora_up"]["weight"].mesh_axes, (None, "data"))
+        self.assertEqual(tuple(param_specs["lora_down"]["weight"].mesh_axes), ("data", None))
+        self.assertEqual(tuple(param_specs["lora_up"]["weight"].mesh_axes), (None, "data"))
 
     def test_forward(self):
         input_dim = 2