fix(test): ignore the serialized Mosaic kernel when diffing HLO

[ecnal-cienet]

Description

The three tests/integration/hlo_diff_test.py cases fail on every scheduled image build (run 29261172318), and have since the stable-image test job was added in #4401 (052fda676). The compiled graph is not deviating — the test is comparing the install path of jax.

Splash attention is a Pallas kernel, so XLA serializes its Mosaic module into the custom call and dumps it as base64 under backend_config.custom_call_config.body. That payload carries the kernel's MLIR location info, which holds the absolute path of splash_attention_kernel.py — and that path depends on where jax was installed:

environment	path baked into the HLO
wheel + venv (presubmit, where the reference is generated)	/__w/maxtext/maxtext/.venv/lib/python3.12/site-packages/jax/...
Docker image (system python, tests run from /deps)	/usr/local/lib/python3.12/site-packages/jax/...

Decoding the payload out of the failing job's log shows that single string is the entire deviation. Each of the three references differs from its dump by exactly one line, and mask the base64 body and all three are byte-identical. Same jax, same libtpu, same graph — different directory.

filter_line already exists to normalize this class of environment noise (it zeroes stack_frame_id and drops the sharding/metadata section lines), so the Mosaic payload is normalized there too. The kernel is still covered by the rest of the dump: its operand shapes and layouts are plain text, and its block sizes are in the plain-text xprof_metadata on the neighboring line.

The reference files do not need regenerating. filter_line runs over the fresh dump and the stored reference alike at compare time, so the committed references — runner path and all — normalize to the same masked form. This is also why running the Update HLO References workflow does nothing: it regenerates in the venv environment, which already matches, and the commit step then fails with "nothing to commit".

The --ignore=tests/integration/hlo_diff_test.py carried for the jax-nightly image is left alone. That image runs a different jax and XLA, so it may deviate for reasons that are real; worth revisiting separately once the stable image is green.

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!