Fix AotHloIdenticalTest false-positive failure on JAX nightly due to new source mapping format

[darisoy]

Description

This PR updates the Ahead-Of-Time (AOT) identical graph verification tests to ignore new XLA debug source location metadata. This fixes false-positive integration test failures occurring under newer JAX/XLA nightly environments.

The Problem Being Solved:

JAX/XLA nightly recently introduced a new debug source mapping metadata format in the compiled HLO headers (file_location_id and parent_frame_id). Because Ahead-Of-Time compilation (train_compile.py) and normal JIT compilation (train.py) have different call stacks, these generated frame IDs differ.

The test's sanitation regex location_map_pattern did not filter out file_location_id, triggering a false-positive AotHloIdenticalTest::test_default_hlo_match failure even though the actual compiled hardware instructions are 100% identical.

The Solution:

Updated the location_map_pattern regex in tests/integration/aot_identical_test.py to match and strip both file_name_id and the new file_location_id metadata blocks before comparing AOT and Real hashes:

location_map_pattern = re.compile(r"^\s*\d+\s+\{(file_name_id|file_location_id)=[^}]*\}\s*$")

FIXES: b/514475466
FIXES: #3941

---

Tests

Build Images workflow no longer fails with AOT identical graph test.

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!