Change in axlearn/common/flash_attention/layer_test.py to use "as_partition_spec" function

to convert Tuple to PartitionSpec.

Author: QuinnMMcGarry

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, jax.P(test_layer.config.mha_dim_to_partition_spec["btnh"][:2]))
+            self.assertEqual(spec, as_partition_spec(test_layer.config.mha_dim_to_partition_spec["btnh"][:2]))
 
     @parameterized.product(
         _TEST_CONFIGS,