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/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"