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)