feat: enable JAX persistent compilation cache by default (+ fix XLA_FLAGS clobber)

[Jammy2211]

Overview

Roll out the verdict of autolens_profiling#71: enable the JAX persistent compilation cache by default, so the prohibitive XLA compile times (117 s → 2.3 s local; 5518 s → 937 s A100 end-to-end when warm) become a once-per-machine cost. The natural home is autoconf/jax_wrapper.py, which already manages JAX env defaults before any JAX import (the x64 pattern).

Included bug fix (same lines): jax_wrapper.py currently overwrites XLA_FLAGS when --xla_disable_hlo_passes=constant_folding is absent, silently discarding user/job flags (--xla_dump_to, --xla_gpu_autotune_level=0). This made --xla_dump_to appear "inert" during #71 and means the historical 2026-07-15 "autotune ruled out" A/B never actually flipped autotune (both runs identical because both were clobbered) — that claim is unproven pending a re-test.

Plan

• Fix the XLA_FLAGS clobber: append --xla_disable_hlo_passes=constant_folding to existing flags instead of replacing.
• Cache-by-default via env (respecting pre-set values, including empty = disable): JAX_COMPILATION_CACHE_DIR → ~/.cache/pyauto_jax (XDG_CACHE_HOME-aware); JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS → 1.
• First-compile UX log line (cache location + one-time-cost explanation).
• test_autoconf/ unit tests for the env handling (reload-based, no JAX import).
• Ship validation: empirically confirm jax.config picks up the env vars (the support enum in fitsable #71 probe set env AND config; the env-only path must be proven, not assumed).
• Ripple corrections: amend autolens_profiling/jax_compile/README.md + support enum in fitsable #71 ("dump inert" → "clobbered by the wrapper"; autotune claim downgraded to unproven).

Detailed implementation plan

Affected Repositories

• PyAutoConf (primary — autoconf/jax_wrapper.py, test_autoconf/)
• autolens_profiling (correction commit to the support enum in fitsable #71 research note, on its existing PR features/jax_ndarray_fitsable #73 branch or follow-up)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoConf	main	clean

Suggested branch: feature/jax-cache-default

Implementation Steps

1. autoconf/jax_wrapper.py: rework the XLA_FLAGS block — parse existing value, append the constant_folding pass-disable if missing (preserve everything else); keep the informational log.
2. Same file: cache block after the x64 block — os.environ.setdefault("JAX_COMPILATION_CACHE_DIR", <default>) guarded by "unset only" (explicit empty string disables); same for JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS = "1"; default dir $XDG_CACHE_HOME/pyauto_jax or ~/.cache/pyauto_jax; log line.
3. test_autoconf/test_jax_wrapper.py: reload-based tests — flags appended not clobbered; cache dir set when unset; pre-set values (incl. empty) respected; no import jax anywhere.
4. Downstream test suites at ship (autoconf is the base layer).

Key Files

• autoconf/jax_wrapper.py — the change
• autolens_profiling/jax_compile/README.md — correction target
• Evidence: autolens_profiling#71 (verdict + measurements), PR features/jax_ndarray_fitsable #73

Original Prompt

PyAutoMind draft/feature/autofit/enable_the_jax_persistent_compilation_cache_by.md (target corrected PyAutoFit → PyAutoConf at start_dev: the wrapper lives in autoconf and env-based defaults remove the workspace-config leg).

[Jammy2211]

Shipped: PR #128 (pending-release). Tests: 147 autoconf + 1493 PyAutoFit downstream against the branch. Env-only path empirically verified (jax.config picks up both vars; cache entries written in fresh processes). Heart RED at ship — 5 pre-existing unrelated reasons, human-acked for PR-open; merge stays human. Ripple corrections to autolens_profiling#71/PR#73 posted (dump-flags clobber, autotune claim downgraded to unproven).

[Jammy2211]

Merged via PR #128 — JAX persistent compilation cache now on by default; XLA_FLAGS clobber fixed. Closing.