From 188af874fb7d28b4ae9a344d4bb5a0f46aaaddba Mon Sep 17 00:00:00 2001 From: Nitin Gangahar Date: Tue, 28 Jul 2026 12:42:24 -0700 Subject: [PATCH] Fix logical to physical axis mapping for attention heads. PiperOrigin-RevId: 955411498 --- 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 2733e77810..6e93c825e1 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -1506,7 +1506,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)