Only pass in logical axis rules if present in config.

Kevin Wang · Google-ML-Automation · commit 3ac667abec01 · 2026-04-10T22:18:47.000-07:00
PiperOrigin-RevId: 898044030
diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py
@@ -169,7 +169,7 @@ def __call__(self, inputs: Array, model_mode: str = MODEL_MODE_TRAIN) -> Array:
             "activation_embed",
         )
     )
-    out_pspec = logical_to_mesh_axes(output_axis_names, self.mesh, rules=self.config.logical_axis_rules)
+    out_pspec = logical_to_mesh_axes(output_axis_names, self.mesh, rules=getattr(self.config, "logical_axis_rules", None))
 
     out_sharding = NamedSharding(self.mesh, out_pspec) if self.config.shard_mode == ShardMode.EXPLICIT else None
 

Original file line number	Diff line number	Diff line change
`@@ -169,7 +169,7 @@ def __call__(self, inputs: Array, model_mode: str = MODEL_MODE_TRAIN) -> Array:`
`169`	`169`	`"activation_embed",`
`170`	`170`	`)`
`171`	`171`	`)`
`172`		`- out_pspec = logical_to_mesh_axes(output_axis_names, self.mesh, rules=self.config.logical_axis_rules)`
	`172`	`+ out_pspec = logical_to_mesh_axes(output_axis_names, self.mesh, rules=getattr(self.config, "logical_axis_rules", None))`
`173`	`173`
`174`	`174`	`out_sharding = NamedSharding(self.mesh, out_pspec) if self.config.shard_mode == ShardMode.EXPLICIT else None`
`175`	`175`