Make the persistent compilation cache actually cache pylcm's programs by hmgaudecker · Pull Request #397 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-07-03T05:53:35Z

Closes #392.

Diagnosis

The persistent compilation cache was never populated, so every fresh process recompiled the whole model. The cause is not an AOT/cache incompatibility: JAX only writes an executable to the cache when its compile time exceeds jax_persistent_cache_min_compile_time_secs, which defaults to 1 second. A pylcm model compiles as many small per-regime/per-period programs, and nearly all of them compile in well under a second — so with JAX's default threshold the cache directory stays empty and nothing is ever reused.

Once the cache is populated, the AOT-compiled programs (.lower().compile() in both solve and simulate) hit it just like lazily jitted ones. On the mortality example (CPU and GPU alike), a fresh-process re-run drops from ~1.7s to ~0.5s in simulate with ~70 cache entries written; the residual is Python-side tracing and lowering, which no XLA cache can skip.

What this PR does

Importing lcm zeroes the threshold:

if os.environ.get("JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS") is None:
    jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)

Applied via jax.config (not os.environ.setdefault) so it takes effect even when jax is imported before lcm — JAX reads the environment variable only once, at import time.
A user-set JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS wins.

This makes the third JAX default pylcm sets at import, next to XLA_PYTHON_CLIENT_PREALLOCATE and JAX_COMPILATION_CACHE_DIR. The user guide (installation and tuning pages) documents it.

Tests

Fresh-interpreter subprocess tests in tests/test_compilation_cache.py: the threshold is zero after import jax; import lcm (the adversarial import order); a user-set threshold survives the import; and — the actual #392 regression — compiling a sub-second kernel writes an entry to a fresh cache directory.

Base branch

Targets feat/terminal-rows (#396); lands on main via the usual cascade. The change is independent of that branch's content — it just continues the stack.

🤖 Generated with Claude Code

…river) Two tests on the 4-CPU distributed seam: - test_distributed_solve_matches_single_device_per_type: the F6 correctness contract — a sharded solve must equal the single-device solve per type slice. Green now; the co-map fix must preserve it (catches a wrong device-local index). - test_distributed_solve_kernel_does_not_all_gather_continuation_v: the red driver — the backward-induction kernel currently all-gathers the continuation V across the type shard; it must read only its device-local slice. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

vmap_1d gains co_mapped_in_axes: an optional per-argument in_axes override so a pytree argument's leading axis can be mapped in lockstep with the mapped variables. The backward-induction co-map uses it to slice each next_regime_to_V_arr leaf to the device-local type, so the continuation-V interpolation reads only its own shard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Fixed, distributed states (e.g. a permanent type sharded one block per device) never transition, so a regime's continuation value depends only on its own slice of the next-period V-array. The grid-search solve kernel now co-maps each such state with the matching axis of every next_regime_to_V_arr leaf that carries it: an outer vmap peels the leading axis off both the state grid and the continuation V, so the interpolation reads only the device-local slice and XLA inserts no all-gather of the full V-array onto every device. - max_Q_over_a splits the co-mapped (leading) states from the inner productmap and wraps it in per-state co-map vmaps; the V-interpolator drops those coordinates and Q_and_F omits the sliced next-states. - processing detects the co-mappable states (distributed and identity-transition) and builds per-state, per-leaf in_axes so a target regime that prunes the state keeps its full leaf. - Simulation is unchanged: it keeps the full continuation V (subjects are not type-aligned across devices), so the co-map applies to the solve path only. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…nuation-v

…seam (#385) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>