Fix logical to physical axis mapping for attention heads.

[copybara-service]

Fix logical to physical axis mapping for attention heads.

Description

Use self._logical_to_mesh_axes to pass logical rules to logical_to_mesh_axes, ensuring the logical axis (e.g. 'tensor') is correctly mapped to the physical mesh axis. This fixes crashes when running with logical sharding configs.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'MaxText Test Execution Times'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 1.15.

Benchmark suite	Current: 9daa835	Previous: df95a1f	Ratio
Total CPU-TORCH-REFERENCE Tests Duration	288.14900000000006 sec	226.67700000000002 sec	1.27

This comment was automatically generated by workflow using github-action-benchmark.