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.

camiloCienet · camiloCienet · commit bda8d78ebfaa · 2025-11-11T05:27:58.000Z
diff --git a/axlearn/common/state_builder_test.py b/axlearn/common/state_builder_test.py
@@ -293,6 +293,7 @@ def _mock_bert_trainer_config(
                 embedding_cfg=bert.bert_embedding_config(max_position_embeddings=max_len),
                 stack_cfg=bert.bert_transformer_config(num_layers=3, num_heads=2),
             ),
+            mesh_axis_names=("fsdp", "model")
         )
 
     def _mock_bert_trainer_config_and_state(

Original file line number	Diff line number	Diff line change
`@@ -293,6 +293,7 @@ def _mock_bert_trainer_config(`
`293`	`293`	`embedding_cfg=bert.bert_embedding_config(max_position_embeddings=max_len),`
`294`	`294`	`stack_cfg=bert.bert_transformer_config(num_layers=3, num_heads=2),`
`295`	`295`	`),`
	`296`	`+ mesh_axis_names=("fsdp", "model")`
`296`	`297`	`)`
`297`	`298`
`298`	`299`	`def _mock_bert_trainer_config_and_state(`