Disable JAX compilation cache when dump_hlo is True

[igorts-git]

Description

When dump_hlo=true is set, MaxText relies on XLA compiler flags to dump HLO files during compilation. However, if jax_cache_dir is configured, subsequent runs reuse the JAX compilation cache on disk. Since compilation is bypassed, XLA does not run, and no HLO files are dumped. This silently fails to produce HLOs if the user clears the dump directory between runs.

Solution: Automatically bypass the JAX compilation cache initialization in pyconfig.py if dump_hlo is True. This forces recompilation and ensures that HLOs are always generated as requested. A warning is logged to inform the user:

JAX compilation cache is disabled because dump_hlo is True. HLO dumping requires recompilation.

Alternative considered: We explored programmatically extracting the final HLO text from the cached executable using compiled.as_text(). While this successfully retrieves the post-optimization HLO even on cache hits, it is unable to recover the full suite of XLA compiler intermediate output files (like before_optimizations.txt, tpu_comp_env.txt, memory usage reports, etc.) which are never generated when compilation is skipped. Forcing recompilation remains the only way to guarantee the complete dump.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/440500124

Tests

Tested on my TPU VM using a tiny model before and after the fix to confirm that when dump_hlo==True XLA compilation gets triggered and HLO is dumped.

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.

[gobbleturk]

This PR implies any hero run should not dump HLO, which is fine I guess since our main HLO use cases are for one-off debugging, but something to note in PR description perhaps, maybe also in types.py documentation

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!