From 7e3d8bb998063dd33a973b4a74302f7871e6a744 Mon Sep 17 00:00:00 2001 From: Nitin Gangahar Date: Fri, 24 Jul 2026 14:38:24 -0700 Subject: [PATCH] Fix logical to physical axis mapping for attention heads. PiperOrigin-RevId: 953561420 --- src/maxtext/layers/attention_op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index cf076b4794..bec4d10f1b 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -1505,7 +1505,7 @@ def wrap_jax_splash_kernel(multi_head_mask, shard_head_size=1): ) return splash_kernel - head_physical_axes = logical_to_mesh_axes((HEAD,), self.mesh)[0] + head_physical_axes = self._logical_to_mesh_axes((HEAD,))[0] head_physical_axes = (head_physical_axes,) if isinstance(head_physical_axes, str) else (head_physical_axes or ()) shard_head_size = math.prod(self.mesh.shape.get(ax, 1) for ax in head_physical_axes) splash_kernel = wrap_jax_splash_kernel(multi_head_mask, shard_head_size)