Fixes 13 unit tests related with mesh_axis naming from CPU Unit tests Jax 0.8.0 and Jax 0.8.1.dev.. with Python 3.12.

Author: camiloCienet

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"