fix: jax_grad assertions fail locally but pass in CI

[Jammy2211]

Overview

Running the jax_grad scripts locally under the resolved smoke profile produces
deterministic assertion failures in scripts that pass in CI on the same commit.
imaging/jax_grad/lp.py is the decisive control: it passes in CI on two independent
runs and fails locally, so the three local failures found during PyAutoHands#226 were
environment artefacts, not source defects.

This is an active trap for anyone validating these scripts locally — during #226 it
looked exactly like two fresh correctness regressions on current main. Whatever the
cause, either the scripts or the documented local-run recipe should make this
reproducible, so that a local FAIL means something.

Split out of PyAutoHands#226; see PyAutoMind/complete/2026/08/jax-grad-smoke-timeout-budget.md.

Plan

• Capture both environments as data before anything drifts. pip freeze +
  np.show_config() from the local venv, and the same from a CI job. The prompt's
  table is currently the only record of the local side.
• Run the decisive A/B against the control. Build the CI install set, confirm
  lp.py PASSES, then change only numpy 2.4.6 -> 2.2.6 and confirm it FAILS. Run
  it in both directions — the upgrade recovering PASS is what rules out
  install-ordering effects.
• If numpy is not the discriminator, diff the two freezes and bisect the
  remaining delta against lp.py, cheapest-first.
• Root-cause the mechanism, not the correlation. lp.py's failure is an
  active-set change in the positive-only NNLS solve, not a last-digit tolerance
  breach — that points at the BLAS/LAPACK backend the numpy wheel ships, not at
  numpy's Python layer.
• Then decide the verdict — either the assertions are under-specified across the
  supported dependency range, or the local environment is out of spec. These call for
  different fixes, and only measurement chooses between them.
• Make a local FAIL mean something: whichever verdict, land a documented,
  reproducible local-run recipe in AGENTS.md and cross-reference it from the
  existing profile_smoke.yaml warning.

What the code already tells us

• scripts/misc/util.py:185 — assert_eager_jit_consistent(..., rtol=1e-10) is
  documented as a pure_callback constant-folding detector, not a numerics
  check. Constant-folding produces order-unity disagreement; the observed 8.9e-6 is
  ~5 orders of magnitude too small to be the thing this assert exists to catch.
• scripts/imaging/jax_grad/lp.py:213 — the failing guard is
  any(|ad[source]| > 1e-3), i.e. the NNLS solve zeroed the source. The linear
  solve is landing on a different active set — the same signature as the
  pixelization.py 8.9e-6 gap and the regularization.py marginal FD breach.
• .github/scripts/smoke_install.sh pins nothing but jax<0.7/jaxlib<0.7 — numpy
  is whatever the resolver picks that day. "The CI set" is not a fixed target, and
  nothing in the repo tells a local runner how to reproduce it.
• AGENTS.md documents run_smoke.py and run_all_scripts.sh but has no
  local-environment recipe at all. That gap is the trap.
• config/build/profile_smoke.yaml:74 already carries a "Do NOT re-derive these from
  local runs" warning naming this exact divergence — the right place to record the
  resolved cause.

Detailed implementation plan

Work Classification

Workspace (single repo) — routes to /start_workspace.

Affected Repositories

• autolens_workspace_test (primary)

Branch Survey

Repository	Current Branch	Dirty?
PyAutoLabs/autolens_workspace_test	main @ b18fd38	clean

worktree_check_conflict jax-grad-local-vs-ci-assertions autolens_workspace_test
-> exit 0, no conflict.

Warning (not a block): feature/jax-grad-param9-mismatch and
claude/rectangular-mesh-gradients-mh1j0z exist on the remote as unregistered
prior jax_grad work.

Suggested branch: feature/jax-grad-local-vs-ci-assertions

Worktree root: ~/Code/PyAutoLabs-wt/jax-grad-local-vs-ci-assertions/
(created later by /start_workspace)

Implementation Steps

1. Capture ground truth. From the local machine: pip freeze > local.txt and
   python -c "import numpy; numpy.show_config()". From CI: the same two, pulled
   from a fresh main smoke run's job log, or reproduced by running
   .github/scripts/smoke_install.sh against fresh chain checkouts. Record jax,
   jaxlib, numpy, scipy and tfp-nightly explicitly.

2. The A/B. In an isolated venv holding the CI set, run
   scripts/imaging/jax_grad/lp.py via env_config.build_env_for_script with the
   workspace root as cwd (the exact env the report used) -> expect PASS.
   pip install numpy==2.2.6 alone -> rerun -> expect the
   All source-parameter gradients are ~zero failure. Then pip install numpy==2.4.6
   -> rerun -> expect PASS to return. lp.py is the control because it is the only
   script known to pass in CI and fail locally, and at ~41s it is the cheapest.

3. Widen if needed. diff local.txt ci.txt; bisect the remaining delta
   package-by-package against lp.py.

4. Mechanism. Compare np.show_config() across the two numpy versions (OpenBLAS
   build and threading). Check whether the inversion's non-negative solve (the fnnls
   active set) selects a different set under each, by printing the active-set size at
   the lp.py evaluation point in both. This is what distinguishes "float noise"
   from "a different solution" — the evidence says it is the latter.

5. The fix, chosen by the measurement:

   • If the supported dependency range genuinely admits this spread: the rtol=1e-10
     on assert_eager_jit_consistent is the outlier and the only tolerance worth
     changing — raised to a value that still detects constant-folding (which is
     order-unity), with the measured spread written in as its justification. Leave
     assert_gradients_match (rtol=1e-3/atol=1e-4) alone unless
     regularization.py's breach survives step 4.
   • If the local env is out of spec: no tolerance change at all. Add a
     "Local runs of jax_grad/" subsection to AGENTS.md with the exact install
     (mirroring smoke_install.sh plus the resolved pins), the build_env_for_script
     invocation, and the cwd requirement.
   • Either way: cross-reference the verdict from the
     config/build/profile_smoke.yaml "Do NOT re-derive these from local runs" note,
     so that warning points at a resolved cause instead of an open mystery.

6. Validate. Run all 7 live jax_grad scripts under the fixed recipe and confirm
   they agree with CI. Ship as a PR; the per-PR smoke gate is the CI-side confirmation.

Explicitly out of bounds

Any of these masks the trap rather than removing it, and AGENTS.md already states
the rule (never edit a script to mask a real regression):

• moving lp.py's evaluation point;
• adding skip_indices to get green;
• widening a tolerance without the measured basis.

Key Files

• scripts/misc/util.py — assert_eager_jit_consistent (:185),
  assert_gradients_match (:157), fd_gradient (:55). Blast radius: all 8
  jax_grad scripts plus scripts/weak/jax_grad.py.
• scripts/imaging/jax_grad/lp.py — the control; failing guard at :213.
• scripts/imaging/jax_grad/pixelization.py, scripts/imaging/jax_grad/regularization.py
  — the two other local failures.
• .github/scripts/smoke_install.sh — the CI install set (unpinned numpy).
• config/build/profile_smoke.yaml:74 — the existing warning to cross-reference.
• AGENTS.md — "Running Scripts" / "Testing"; where the local recipe lands.

Trade-off

A util.py tolerance change has repo-wide blast radius and needs the strongest
evidence; an AGENTS.md recipe has zero blast radius but only helps people who read
it. Default to documentation unless step 4 proves a real supported-range problem.

Open dependency

The local venv that produced the failures cannot be inspected from a cloud session.
If it still exists, capture its pip freeze + np.show_config() before it drifts —
otherwise step 2 has no ground truth to reproduce.

Sizing note

Brain sized this medium (declared) vs large (derived); the Bug Agent classifies it
severity=critical, scope=single-repo, type=wrong-result, confidence=low, strategy
investigate-first. Reproduce and confirm root cause before patching anything.

Original Prompt

Click to expand starting prompt

jax_grad scripts fail assertions locally that PASS in CI

Type: bug
Target: autolens_workspace_test
Repos:

• autolens_workspace_test
  Difficulty: medium
  Autonomy: supervised
  Priority: medium
  Status: formalised

Running the jax_grad scripts locally under the resolved smoke profile produces
deterministic assertion failures in scripts that pass in CI on the same commit.
Found while measuring script durations for PyAutoHands#226.

Evidence

Run via env_config.build_env_for_script with the workspace root as CWD (i.e. the
exact env the runner builds — PYAUTO_SMALL_DATASETS unset, PYAUTO_DISABLE_JAX
unset, PYAUTO_TEST_MODE=2, verified by printing the resolved env):

script	local	CI (run 30858578587 / 30790463134)
imaging/jax_grad/lp.py	FAIL 41.3s	PASS 39.6s / 40.0s
imaging/jax_grad/knn.py	PASS 141.6s	PASS 200.0s / 175.8s
imaging/jax_grad/pixelization.py	FAIL 57.5s	PASS 244.8s (06:31Z)
imaging/jax_grad/regularization.py	FAIL 131.5s	(import gap, then TIMEOUT)
point_source/jax_grad/gradient.py	PASS 665.9s	TIMEOUT (300s cap)

lp.py is the decisive case: it passes in CI on both runs and fails locally.

Failures are deterministic and bit-identical across repeated runs, e.g.
pixelization.py:

AssertionError: Eager (-8354.484097835004) and jitted (-8354.55843260181) evaluations
disagree — possible pure_callback constant-folding; do not trust jitted gradients.

(relative difference ~8.9e-6 against assert_eager_jit_consistent's rtol=1e-10).

lp.py fails with All source-parameter gradients are ~zero — NNLS zeroed the source;
regularization.py with an AD-vs-FD mismatch marginally over tolerance
(abs_err=[0.045, 0.042, 0.057] vs tolerance=[0.031, 0.008, 0.003]).

What is ruled out

• Not the small-datasets cap. full_datasets correctly unsets
  PYAUTO_SMALL_DATASETS; verified by resolving the env directly rather than
  inferring from mask sizes.
• Not a JAX version difference. Local jax/jaxlib are 0.10.2 — identical to CI.
• Not flake. Repeated runs give bit-identical values.

Prime remaining suspect: numpy 2.2.6 local vs 2.4.6 in CI, or another local venv
package differing from the CI install set. Not yet confirmed.

Why it matters

This is an active trap for anyone validating these scripts locally. During #226 it
looked exactly like two fresh correctness regressions on current main
(pure_callback constant-folding, and an FD tolerance breach). Only running a
control — lp.py, known-passing in CI — revealed that the local environment
itself produces the failures, so none of the three local failures were evidence of
source defects.

Whatever the cause, either the scripts or the documented local-run recipe should make
this reproducible, so a local FAIL means something.

Suggested scope

1. Bisect the local-vs-CI package delta (start with numpy 2.2.6 -> 2.4.6) against
   lp.py, the cleanest discriminator.
2. If numpy is the cause, decide whether the tolerances are under-specified for the
   supported numpy range, or the local env should be pinned to the CI set.
3. Record the outcome in the workspace's local-run instructions.

[Jammy2211]

Root cause found and reproduced. It is not numpy — it is a stale on-disk dataset.

Investigated in a cloud session against a CI-equivalent install (Python 3.12.3, smoke_install.sh replayed against current main of all five libraries, jax/jaxlib 0.10.2, scipy 1.17.1). The laptop that produced the original report was not available, so everything below is reproduced from a clean checkout.

The prime suspect is falsified

imaging/jax_grad/lp.py — the control — run under the resolved smoke profile at three numpy versions:

numpy	result
2.2.6 (the reported local version)	PASS
2.4.6 (the reported CI version)	PASS
2.5.2 (today's fresh resolve)	PASS

The three runs are byte-identical apart from log timestamps — both log-likelihoods (-550210.563030, -48960.221408) and all 40 gradient-table entries reproduce exactly. Pinning to a single core (taskset -c 0) is also byte-identical, which rules out thread-count-dependent reduction order in the XLA CPU backend.

numpy was never mechanically able to cause this: the likelihood and its gradient are computed through JAX/XLA, and numpy only does the finite-difference bookkeeping and the assertion arithmetic afterwards. The 2.2.6-vs-2.4.6 delta was a coincidence in a pip freeze diff.

The actual cause

PyAutoArray/autoarray/util/dataset_util.py:72:

def should_simulate(dataset_path):
    if os.environ.get("PYAUTO_SMALL_DATASETS") == "1":
        if Path(dataset_path).exists():
            shutil.rmtree(dataset_path)
    return not Path(dataset_path).exists()

The guard is existence-only, and asymmetric. It force-regenerates when PYAUTO_SMALL_DATASETS=1, but has no corresponding check in the other direction. Combined with two other facts:

• dataset/** is gitignored (.gitignore:9), so the dataset is never in the checkout;
• every jax_grad script declares ENV: jax full_datasets, which unsets PYAUTO_SMALL_DATASETS.

the consequence is:

• CI clones fresh every run, so dataset/imaging/jax_test/ is always absent and always simulated at full resolution. CI cannot ever hit this.
• Locally the directory persists indefinitely. Once written it is never regenerated under full_datasets, no matter how stale it is.

And the trap: PYAUTO_SMALL_DATASETS=1 is the smoke default for every other script in the repo. So any prior local smoke run rmtree'd the full dataset and rewrote it at reduced resolution. The next jax_grad run then finds a directory that exists, skips simulation, and silently loads 15x15 small-datasets FITS files under full-resolution settings.

This is why "not the small-datasets cap" was ruled out incorrectly. That check verified the resolved env, which was genuinely correct — full_datasets does unset the variable. The damage was done by an earlier run and is baked into the FITS files on disk. No amount of env-resolution checking can see it.

Reproduction (from a clean checkout, deterministic)

step	state	result
1	fresh full dataset, data.fits = 262080 B	lp.py PASS
2	one run with PYAUTO_SMALL_DATASETS=1 -> data.fits = 5760 B	that run PASSES (self-consistent)
3	next run under full_datasets, stale 5760 B dataset reused	all three reported failures reproduce
4	rm -rf dataset/imaging/jax_test -> fresh simulate, 262080 B	lp.py PASS

Step 3 output, against the original report:

• lp.py — AssertionError: All source-parameter gradients are ~zero — NNLS zeroed the source; move the evaluation point. Verbatim match. Source gradients sit at ~1e-12 while lens gradients are ~1e8. stderr carries the giveaway: UserWarning: Mask padded from (16, 16) to (36, 36).
• pixelization.py — Eager (-8354.484097833672) and jitted (-8354.558432600481) evaluations disagree. The report has -8354.484097835004 / -8354.55843260181 — agreement to 11 significant figures (the last digits differ from ordinary library drift since 2026-08-04).
• regularization.py — AD-vs-FD mismatch whose tolerance vector matches the report exactly: reported tolerance=[0.031, 0.008, 0.003], reproduced tolerance=[0.031, 0.008, 0.003, 0.127, 0.042].

What this means for the assertions

No tolerance should be changed. All three assertions did exactly their job — they detected a genuinely invalid dataset. In particular assert_eager_jit_consistent's rtol=1e-10, which looked indefensibly tight, is vindicated: it is a pure_callback constant-folding detector and it correctly refused to certify gradients computed on a corrupt input. Widening it would have deleted a working alarm and left the real defect in place.

The three items listed as out-of-bounds in the plan above (moving lp.py's evaluation point, adding skip_indices, widening tolerances) would each have masked this.

Where the fix belongs

Upstream in PyAutoArray, not in this workspace. should_simulate cannot currently distinguish a dataset simulated under one resolution regime from one simulated under the other, and every PyAuto workspace using the auto-simulate pattern inherits the same trap — this is not specific to jax_grad or to this repo.

Two candidate shapes, with a real trade-off:

1. Regime marker. Stamp the regime the dataset was simulated under and regenerate when the current regime differs. Correct and cheap at run time, but needs care about what happens when simulation fails part-way (a stamped-but-empty directory would be a worse failure mode than today's).
2. Symmetric regeneration. Make the full path rmtree the same way the small path already does. Trivially correct and a 2-line change, but costs a re-simulation on every local run.

Handing the choice to the library owner rather than picking unilaterally — it changes this task's classification from workspace to library.

Incidental finding (separate defect, not this bug)

.github/scripts/smoke_install.sh line 9 is stale:

pip install "jax<0.7" "jaxlib<0.7"

This downgrades jax to 0.6.2 and raises a resolver conflict against autonerves' own jax<0.11.0,>=0.7.0 requirement. The install only ends up on the intended 0.10.2 because the next line's [optional] extras happen to pull it back up. CI is currently getting the right jax by accident, and any reordering or pin change in the extras would silently drop the whole smoke suite onto jax 0.6.2. Worth its own issue.

---

Generated by Claude Code

[Jammy2211]

Fix implemented and verified — PyAutoArray, not this repo

Pushed to PyAutoArray branch claude/jax-grad-assertions-ci-hontn3, commit 5ef6eea. No workspace change and no tolerance change — the three assertions were correct all along and the fix is entirely upstream.

The design went through an adversarial review before implementation, which corrected three things in the original proposal (recorded here because two of them would have introduced a worse bug than the one being fixed):

• The pixel-scale check was unimplementable. The FITS headers carry only SIMPLE / BITPIX / NAXIS / NAXIS1 / NAXIS2 — pixel scale is supplied by the caller at load time and never written. Dropped.
• The predicate is == (16, 16), not <=. Mask2D.circular caps to exactly the cap shape (mask_2d.py:371-373), so the cap can never emit 12x12. Widening a destructive rmtree predicate past what the cap can produce buys no detection and only risks real data.
• The check keys on data.fits by name, never "the first FITS in the directory". PSF kernels are legitimately tiny at full resolution — 11x11 appears in 49 places in autolens_workspace, and dataset/cluster/test/psf.fits is on disk right now at 5760 B, byte-identical in size to a capped data.fits. A glob-based check would have regenerated every dataset carrying a PSF on every single run: the exact performance problem that ruled out unconditional regeneration, reintroduced through the back door.

"Unknown means no" throughout: a missing, unreadable or non-2D data.fits preserves the existing existence-only behaviour rather than deleting.

Verification

• Control-tested. The small->full regression test returns False against the unfixed body and True against the fixed one — it fails without the fix, so it is testing something.
• 1168 tests pass (full test_autoarray suite). 12 new tests: all four regime transitions, plus false-positive guards that each assert a dataset is preserved (tiny PSF, PSF at exactly the cap shape, 12x12 data, JSON-only dataset, unreadable data.fits).
• End-to-end against the original failure. With the patched library installed: poison the dataset with one PYAUTO_SMALL_DATASETS=1 run (data.fits -> 5760 B), then run imaging/jax_grad/lp.py under the full_datasets profile. Before the fix that raised All source-parameter gradients are ~zero; after it, the dataset is regenerated at 262080 B and the checks pass.
• No steady-state cost. With a full dataset present, should_simulate returns False in 0.64 ms — a header read, no re-simulation. (The first call in a process also pays a ~206 ms astropy.io.fits import that import autolens does not itself trigger, but any script reaching this point loads FITS moments later regardless, so it is paid earlier rather than added.)

Scope — this closes the imaging case only

Filed rather than silently left open:

• feat: stamp the small-datasets regime into FITS headers at the writer funnel PyAutoNerves#153 — stamp the regime into FITS headers at the single writer funnel (fitsable.py:89). This is the only discriminant that can catch the interferometer case, where a capped run writes identical NAXIS with different values and therefore fails silently — no shape mismatch, no assertion. Deliberately not folded into this fix: it changes a header card on every FITS the stack writes, which could disturb round-trip tests and file-hash pins.
• bug: should_simulate rmtree's committed, gitignore-allowlisted datasets under PYAUTO_SMALL_DATASETS=1 PyAutoArray#470 — a separate live defect found on the way: the small-datasets branch rmtrees dataset/point_source/simple, which is committed and explicitly allowlisted in .gitignore:13, replacing it with output from a solver that short-circuits to a fixed position pair under the cap. Recoverable via git checkout, but it violates the invariant the allowlist exists to express, and no variant of this fix touches it.

Point-source and weak-lensing datasets (JSON, no FITS) remain exposed to the underlying bug and regress to existence-only behaviour. That is stated in the should_simulate docstring rather than left for the next person to rediscover.

Incidental, still unfiled

.github/scripts/smoke_install.sh:9's pip install "jax<0.7" "jaxlib<0.7" downgrades jax to 0.6.2 and raises a resolver conflict against autonerves' own jax<0.11.0,>=0.7.0; the install only lands on the intended 0.10.2 because the next line's [optional] extras pull it back up. CI is getting the right jax by accident. Left for a maintainer to file or fold into other install work, since it is unrelated to this bug.

No PR opened — the branch is pushed and ready if you want one.

---

Generated by Claude Code

[Jammy2211]

Closing — fix merged as PyAutoLabs/PyAutoArray#471 -> 2a6a416, with all three checks green (unittest (3.12), unittest (3.13), unittest-nojax).

Completion record: PyAutoMind/complete/2026/08/jax-grad-local-vs-ci-assertions.md.

Closing this because the reported symptom is resolved: a local jax_grad FAIL now means something again. The wider bug class is not closed — point-source and weak-lensing datasets (JSON, no FITS) and interferometer datasets (shape-invariant under the cap, so they fail silently) remain exposed, tracked at PyAutoLabs/PyAutoNerves#153 and PyAutoLabs/PyAutoArray#470. Reopen or reference those rather than this one if a stale-dataset symptom shows up in a non-imaging family.

---

Generated by Claude Code