Merge pull request Borklet-Labs#6 from CIeNET-International/jax_nightly_py3.12_test

andersensam · web-flow · commit 23633d81d45d · 2025-11-11T11:22:21.000-08:00
Fixes 13 unit tests related with mesh_axis naming from  CPU Unit test…
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(`