feat: JAX persistent compilation cache by default + XLA_FLAGS append fix (#127) - #128
Merged
Conversation
…AGS clobber JAX_COMPILATION_CACHE_DIR defaults to ~/.cache/pyauto_jax (XDG-aware) when unset; empty string disables; JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS defaults to 1. XLA_FLAGS now appends constant_folding disable instead of overwriting user/job flags. Evidence: autolens_profiling#71 (cache 51x local / 5.9x A100 end-to-end). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Rollout of the autolens_profiling#71 verdict (closes #127): enable the JAX persistent compilation cache by default, turning prohibitive XLA compile times into a once-per-machine cost — measured 117 s → 2.3 s locally (51×) and 5518 s → 937 s end-to-end on the A100 pathological case (5.9×).
Also fixes the
XLA_FLAGSclobber: the wrapper previously overwrote the variable when the constant_folding disable was absent, silently discarding user/job flags (--xla_dump_to,--xla_gpu_autotune_level=0). It now appends. This bug made dump flags look inert during the #71 research and invalidated the historical 2026-07-15 "autotune ruled out" A/B (downgraded to unproven on the #71 record).API Changes
None (public Python API untouched). Environment behavior at
import autoconf.jax_wrapper(i.e.import autofitand downstream):JAX_COMPILATION_CACHE_DIR— newly defaulted to$XDG_CACHE_HOME/pyauto_jaxor~/.cache/pyauto_jaxwhen unset. Pre-set values respected; explicit empty string disables the cache.JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS— defaulted to1when unset and the cache is enabled; pre-set values respected.XLA_FLAGS— existing user flags are now preserved (constant_folding disable appended instead of replacing the whole value).Tests
test_autoconf/test_jax_wrapper.py: 9 env-handling tests (reload-based, no JAX import per the unit-test rule).jax_wrapperat__init__line 1 — the highest-risk downstream of this base-layer change).jax.config.jax_compilation_cache_dir/jax_persistent_cache_min_compile_time_secspick up the values, and cache entries are written on first compile.Heart
RED at ship time — pre-existing reasons unrelated to this branch; human-acknowledged for PR-open (merge stays human).
🤖 Generated with Claude Code