Fixes CPU Unit tests Jax 0.8.0 with Python 3.12

Author: camiloCienet

axlearn/common/state_builder_test.py

@@ -315,7 +315,6 @@ def test_truncation(self, strategy):
             strategy=strategy,
         )
         converter: PosEmbeddingConverter = cfg.instantiate(parent=None)
-
         converted_state = converter.source_to_target(source_state, target_state)
         self.assertEqual(
             (1, 512, 32),

axlearn/common/t5_test.py

@@ -342,7 +342,7 @@ def test_against_t5_encoder_decoder_model(self, arch):
     def test_pjit(self):
         # A simple test to ensure a train step does not leak tracers.
         mesh_shape = (1, 1)
-        mesh_axes = ("data", "model")
+        mesh_axes = ("fsdp", "model")
         devices = mesh_utils.create_device_mesh(mesh_shape)
         with jax.checking_leaks(), jax.sharding.Mesh(devices, mesh_axes):
             vocab_size = 6

axlearn/common/test_utils.py

@@ -727,7 +727,7 @@ def pytest_addoption_atomic(parser, option, **kwargs):
 def mock_trainer_config(
     input_config: InstantiableConfig,
     model_config: BaseModel.Config,
-    mesh_axis_names: Sequence[str] = ("data", "model"),
+    mesh_axis_names: Sequence[str] = ("fsdp", "model"),
 ) -> SpmdTrainer.Config:
     cfg = SpmdTrainer.default_config()
     cfg.name = "mock_trainer_config"