fix: restore zero-contour compatibility with JAX 0.11

[Jammy2211]

Overview

Restore optional zero-contour calculations with JAX 0.11.2 without changing global third-party behavior or importing JAX at library import time. This blocks Euclid pipeline PR #90 and the approved source-checkout deployment.

Plan

• Add a narrow lazy compatibility adapter used by both LensCalc solver construction sites.
• Cover values, compilation, batching behavior and derivatives in workspace integration scripts on JAX 0.10.2 and 0.11.2.
• Validate the NumPy library suite and Euclid regression, obtain independent review, then ship library first through readiness and green-CI gates.

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary; source)
• autogalaxy_workspace_test (integration coverage)

Branch Survey

Both repositories are clean on main; no conflicting claims or existing repair worktrees. PyAutoGalaxy HEAD 33714b8; workspace_test HEAD e8b2f38. Existing workspace branch feature/jax-runtime-and-parity is not checked out or claimed.

Suggested branch: feature/euclid-jax-contour-compat
Work classification: Both, library first.
Worktree root: /home/jammy/Code/PyAutoLabs-wt/euclid-jax-contour-compat

Implementation Steps

1. Add an internal lazy solver factory in autogalaxy/operate/ which subclasses optional jax_zero_contour.ZeroSolver and normalizes the callable in step_parallel_tol with jax.tree_util.Partial. Preserve tolerance/iteration arguments, inherited numerical algorithms and the LensCalc solver cache. No global monkey patch or module-level JAX imports.
2. Use it at both solver construction sites in autogalaxy/operate/lens_calc.py, preserving optional-dependency warnings and return contracts.
3. Add integration coverage under scripts/misc/latent/ without growing the curated smoke list. Check circle contour and Einstein-radius values, jit, gradients and batching against both versions. Respect the existing unsupported-vmap boundary; do not claim expanded support without evidence.
4. Run the full PyAutoGalaxy NumPy suite, affected workspace smoke checks, and test_latent_euclid_variables_traces_under_jax_jit against both JAX versions with unchanged jax-zero-contour 2.0.0. Retain current tests, no skips/xfails or CI downgrade. Independent review must be CLEAN.
5. Ship source then integration PR through normal gates. Rerun pipeline Include configs should apply to every image type #90 CI after upstream merge; merge only when every required check is green.

Authorization and gates

The user explicitly approved this source plan and green-CI repair/pipeline merges in the live 2026-09-18 session. This supersedes the original prompt's proposed-plan status. No library Heart override is inferred from pipeline issue #89. Current Heart RED: release validation FAILED (stage integrate); YELLOW: manifest drift: remote-session blocks (generated) — 2 mismatch(es) vs PyAutoMind/repos.yaml. Follow the applicable task-specific shipping gate after tests/review. No release authorization.

No SLURM submissions, including smoke jobs: “dont submit jobs until Ive okayed modeling script”. Preserve all science data/outputs. No fields.field.shear compatibility or tests; tracer.fields remains a list.

Original Prompt

Starting prompt (historical approval status superseded above)

Restore JAX 0.11 zero-contour compatibility for Euclid deployment

Type: bug
Target: @PyAutoGalaxy @autogalaxy_workspace_test
Autonomy: human-required
Priority: high
Filed: 2026-09-18
Blocks: PyAutoLabs/euclid_strong_lens_modeling_pipeline#90
Status: proposed; source plan approval required

Original user request

Euclid + euclid_dr1 deployment (PyAutoMind/tmp/handoffs/euclid-flat-fields.md)

The user approved the pipeline plan and subsequently approved its task-specific Heart RED development override and merge only with all required CI green. This new upstream source repair is proposed, not yet approved. The science modelling-script submission hold remains.

Reproduction and fix locus

Pipeline #90 at 7fbdbe5 has 230 local tests and 9/9 smoke passes, independent review CLEAN. Both CI unit matrix legs fail the existing latent JAX trace test with JAX/jaxlib 0.11.2 and jax-zero-contour 2.0.0; local JAX/jaxlib is 0.10.2.

A standalone circle-contour calculation with no PyAuto code passes on JAX 0.10.2, fails on 0.11.2 with functools.partial not a valid JAX type. ZeroSolver.newton uses jax.lax.custom_root; its step_parallel_tol puts the returned callable into a lax.while_loop state. A diagnostic-only adapter wrapping that callable with jax.tree_util.Partial passes the contour calculation and a gradient under 0.11.2. No installed package or production source was patched.

Evidence: PyAutoMind/tmp/euclid-flat-fields/{ci-failed.log,zero-contour-repro.py,repro-jax010.log,repro-jax011.log,zero-contour-candidate.py,repro-candidate.log}. Third-party source: CKrawczyk/Jax-Zero-Contour.

Proposed plan

1. Add a narrowly scoped optional-dependency adapter in PyAutoGalaxy, used by both ZeroSolver construction sites in autogalaxy/operate/lens_calc.py. Normalize the callable to a JAX-compatible Partial in a subclass, without globally monkey-patching the third-party class or importing JAX at module import time.
2. Extend JAX integration coverage in autogalaxy_workspace_test/scripts/misc/latent/ (not library NumPy unit tests): contour/Einstein-radius values, jit, vmap and gradients under both JAX 0.10.2 and 0.11.2. Retain all existing pipeline tests; no skips, expected failures, or CI downgrade.
3. Run PyAutoGalaxy's NumPy suite, affected integration/smoke checks, Euclid failing latent trace test on both JAX versions, and independent review. Ship the separate library repair through standard gates, then rerun pipeline Include configs should apply to every image type #90 CI and resume the originally approved source-checkout deployment. Existing Heart RED reasons still require the proper task-specific shipping authorization for the new library task.

Branch survey

PyAutoGalaxy and autogalaxy_workspace_test are clean on main; conflict guard reports no claims for either. Proposed branch: feature/euclid-jax-contour-compat. Proposed worktree: ~/Code/PyAutoLabs-wt/euclid-jax-contour-compat. No issue, source branch, or source edit has been created for this proposed repair.

[Jammy2211]

Progress — 2026-09-18

Implemented and staged in the two feature/euclid-jax-contour-compat worktrees: a lazy cached ZeroSolver subclass and both LensCalc call sites, plus the standalone workspace integration script. No production environment or upstream dependency was patched. No source commit/push/PR yet.

Validation is green:

• PyAutoGalaxy NumPy suite: 1236 passed.
• New integration regression: PASS on JAX 0.10.2 and 0.11.2, both with jax-zero-contour 2.0.0. Analytic circle values, jit, bounded vmap, gradients; both LensCalc paths and cache reuse.
• Previously failing Euclid latent trace test: PASS on both versions.
• Full Euclid suite on JAX 0.11.2: 230 passed.
• Workspace smoke on JAX 0.11.2: 41/41 passed (existing 40 plus regression).
• Independent Sol review: CLEAN; diff checks clean. Review required removing the addition to the curated smoke list; the standalone _jax.py regression remains and was run directly and through the smoke runner. Final smoke_tests.txt is unchanged.

Scope limitation: arbitrary terminated-contour vmap support remains excluded upstream. A separate Einstein-radius reverse-gradient shape failure reproduces on unmodified PyAutoGalaxy/JAX 0.10.2 and is not claimed fixed here.

Shipping blocker

Source-plan and green-CI merge authorization are recorded, but this task has no task-specific Heart RED development override. Current RED reason: release validation FAILED (stage integrate). YELLOW: manifest drift: remote-session blocks (generated) — 2 mismatch(es) vs PyAutoMind/repos.yaml. Library/workspace PR bodies are prepared; shipping waits on the canonical development-only exception. No release or failed-check bypass is authorized.

Deployment holds

Pipeline #90 remains unmerged at 7fbdbe5. Local science remains d53b9ce with all 18 science commits and untracked preparation preserved; all 21 port preimages still match. CSV/submit ordering and 1000 FITS/info pairs reverified. RAL library SHAs remain unchanged; running job 343413 loads the shared canonical stack, so HPCPullPyAuto must wait. Remote editor backups are preserved; Fit cleanliness still needs a completed recheck. No science port, sync, refresh or submission.

Submission hold remains: “dont submit jobs until Ive okayed modeling script”.

Local evidence: PyAutoMind/tmp/euclid-flat-fields/repair-validation.md, repair-review.md, repair-*.log, prepared repair-library-pr.md and repair-workspace-pr.md.

[Jammy2211]

Human-authorized development-only Heart RED override — 2026-09-18

The live user replied “I authorize” to: “Authorize the development-only Heart override for #622’s repair PRs, acknowledging these reasons, and their green-CI merges?”

Exact current RED: release validation FAILED (stage integrate).
Exact current YELLOW: manifest drift: remote-session blocks (generated) — 2 mismatch(es) vs PyAutoMind/repos.yaml.

Passed branch gates: 1236 library tests; both-version JAX integration and Euclid regression; 230 full Euclid tests under JAX 0.11.2; 41/41 workspace smoke scripts; independent Sol review CLEAN.

Scope: commit, push and pending-release PRs for euclid-jax-contour-compat / #622, then merge only with every required GitHub CI check green, library first. Heart remains RED for releases. No release, failed-check bypass, protection override or SLURM submission is authorized. Existing modelling-script hold remains.

[Jammy2211]

Library PR Created: #623 (f3bab3b). Workspace integration: PyAutoLabs/autogalaxy_workspace_test#123 (8bc2713). Both pending-release labels verified. All local gates passed; CI running, merge library first only when every required check is green. No release or SLURM submission.

[Jammy2211]

Shipped: PyAutoGalaxy#623 merged at 90e757d after all 4 CI checks passed; autogalaxy_workspace_test#123 merged at ae45e490 after all 3 CI checks passed. Both branches verified as ancestors of origin/main. Local validation: 1236 library tests, 230 Euclid tests, two-version regression, 41 smoke scripts, independent review CLEAN. Pipeline #90 subsequently passed all 9 checks and merged. Repair complete; science/RAL deployment remains on pipeline issue #89, currently held by shared-stack job 343413. Heart release RED and SLURM submission hold remain.