Repository navigation
fix: Matern-kernel JAX crash — tfp 0.25.0 incompatible with jax 0.10.2 #385
Description
Overview
Any pixelization using Matern-kernel regularization crashes at import on the current release-resolved JAX stack (jax/jaxlib 0.10.2 + tensorflow-probability 0.25.0). Surfaced by the 2026-07-13 release-validation Stage-3 run (PyAutoHeart workspace-validation run_id=29266305445, TestPyPI 2026.7.13.1.dev65501). It ships broken to real JAX users and is not caught by the curated smoke subset — only the full release-fidelity set exercises Matern pixelization.
Reproduced in a clean env with the exact release pins: import tensorflow_probability.substrates.jax (lazily imported by autoarray/inversion/regularization/matern_kernel.py:37, kv_xp) crashes at .../internal/backend/jax/ops.py:681 with:
jax.interpreters.xla.pytype_aval_mappings[onp.ndarray]
AttributeError: module 'jax.interpreters.xla' has no attribute 'pytype_aval_mappings'
pytype_aval_mappings was removed from jax.interpreters.xla in a JAX newer than TFP 0.25.0 targets — so jax 0.10.2 + tfp 0.25.0 is a broken combination.
Resolution (chosen after evaluating all three prompt options)
- (a) lower the jax cap so tfp 0.25.0 works — infeasible:
pytype_aval_mappingsis already gone by jax 0.9.2, andautoconf[jax]floorsjax>=0.7.0(jaxnnls needs modern jax). You'd need jax < ~0.5. - (b) bump tfp — no stable release exists (0.25.0 is the latest on PyPI), but
tfp-nightlyworks: verifiedtfp-nightly==0.26.0.dev20260713computesbessel_kve(0.5,1.0)=1.2533on jax 0.10.2, and installs via an explicit==pin without--pre. ← chosen - (c) drop tfp, compute
kvdirectly —jax.scipy.specialhas nokv/kve, andnuis a freeUniform(0.5, 5.5)prior (continuous — no closed-form shortcut). The only route is vendoring tfp'sbessel_kve(~500 lines + custom gradients). Durable but too large for this release-blocker → filed as a follow-up.
Plan
- Swap
tensorflow-probability==0.25.0→tfp-nightly==0.26.0.dev20260713in PyAutoArray's[jax]extra (import name is unchanged;matern_kernel.pyneeds no code edit). - Fix the
kv_xpdocstring (it claims ajax.scipy.special.kvfast-path that doesn't exist) and point theImportErrormessage attfp-nightly. - Add a JAX Matern-regularization parity assertion to
autogalaxy_workspace_test/scripts/jax_assertions/(JAX guards don't belong intest_autoarray/, which is NumPy-only) so a broken combo is caught by the validation harness, not just full release-fidelity runs. - Verify in the clean repro venv (jax 0.10.2 + tfp-nightly): import +
matern_kernel(xp=jnp)compute succeed. - File a durable follow-up: vendor
bessel_kveand drop tfp entirely.
Detailed implementation plan
Work Classification
Library + Workspace (library-first).
Affected Repositories
- PyAutoLabs/PyAutoArray (primary — dependency + docstring)
- PyAutoLabs/autogalaxy_workspace_test (regression assertion)
Branch Survey
| Repository | Current Branch | Dirty? |
|---|---|---|
| ./PyAutoArray | main | clean |
| ./autogalaxy_workspace_test | main | clean |
Suggested branch: fix/matern-tfp-jax-compat
Worktree root: ~/Code/PyAutoLabs-wt/matern-tfp-jax-compat/
Implementation Steps
PyAutoArray/pyproject.toml([jax]extra, ~line 61):"tensorflow-probability==0.25.0"→"tfp-nightly==0.26.0.dev20260713".PyAutoArray/autoarray/inversion/regularization/matern_kernel.py: correct thekv_xpdocstring andImportErrormessage (no behaviour change).autogalaxy_workspace_test/scripts/jax_assertions/matern_regularization.py(new): build aMaternKernelregularization matrix for a small mesh withxp=npand withxp=jnpunderjax.jit; assert numpy/jax parity (exercises the tfpbessel_kvepath).- Verify in repro venv.
- Follow-up prompt:
PyAutoMind/draft/refactor/autoarray/matern_vendor_bessel_kve.md.
Key Files
PyAutoArray/pyproject.toml— the dependency pin (the actual fix).PyAutoArray/autoarray/inversion/regularization/matern_kernel.py—kv_xp/ docstring.autogalaxy_workspace_test/scripts/jax_assertions/matern_regularization.py— regression.
Autonomy
--auto, effective level supervised (bug work-type cap). Plan written here; run parks at the ship sign-off with a question rather than opening the PR. Merge/close remain human.
Original Prompt
Click to expand starting prompt
Matern-kernel regularization crashes: tensorflow-probability 0.25.0 incompatible with jax 0.10.2
Type: bug
Target: autoarray
Repos:
- PyAutoArray
Difficulty: medium
Autonomy: supervised
Priority: high
Status: formalised
Any pixelization using Matern-kernel regularization crashes at import on the current
resolved stack (jax/jaxlib 0.10.2, tensorflow-probability 0.25.0). Surfaced by the
2026-07-13 release-validation Stage-3 run (PyAutoHeart workspace-validation
run_id=29266305445, TestPyPI 2026.7.13.1.dev65501) — autogalaxy imaging
scripts/imaging/features/pixelization/galaxy_reconstruction.py FAIL, and the same import
almost certainly drives the autolens imaging + autolens/autogalaxy interferometer
pixelization failures in that run.
Root cause — autoarray/inversion/regularization/matern_kernel.py:37 (kv_xp) lazily does
import tensorflow_probability.substrates.jax as tfp; TFP 0.25.0 then executes
tensorflow_probability/.../internal/backend/jax/ops.py:681:
jax.interpreters.xla.pytype_aval_mappings[onp.ndarray]
AttributeError: module 'jax.interpreters.xla' has no attribute 'pytype_aval_mappings'
pytype_aval_mappings was removed from jax.interpreters.xla in a JAX version newer than
what TFP 0.25.0 targets — so jax 0.10.2 + tfp 0.25.0 is a broken combination. This ships
broken to real users on the current stack; it is NOT gated by the curated smoke subset
(only the full release-fidelity set exercises Matern pixelization), so the nightly release
would not catch it. Scoped to the Matern path — pixelization/likelihood_function.py
passed; only the TFP-importing regularization fails.
Resolution options to weigh (do NOT assume one): (a) add/lower a dependency cap so
jax/jaxlib resolve to a version TFP 0.25.0 supports; (b) bump tensorflow-probability
to a release compatible with jax 0.10.x, if one exists; (c) remove the TFP dependency from
matern_kernel.py by computing the modified-Bessel kv term via jax.scipy/scipy
directly (drops a heavy dep). First find where the jax/tfp caps are declared (autoarray vs
autoconf) — see the dependency graph and the parked astropy_cap_bump cap work —
then reproduce on a clean env with the pinned versions before choosing. Add a regression
test that imports/exercises Matern-kernel regularization under JAX.
Activity
Jammy2211 commentedon Jul 13, 2026
⏸️ Parked at ship sign-off — supervised
--auto(needs your go-ahead to open the PR)Implementation is done and verified; under the
supervisedbug cap an autonomous run stops here for your review of the changes before commit/push/PR. Nothing has been committed or pushed. Merge/close stay human regardless.What changed (worktree
~/Code/PyAutoLabs-wt/matern-tfp-jax-compat, branchfeature/matern-tfp-jax-compat)PyAutoArray (2 files)
pyproject.toml—optionalextra:tensorflow-probability==0.25.0→tfp-nightly==0.26.0.dev20260713(+ a comment explaining why). This is the fix.autoarray/inversion/regularization/matern_kernel.py—kv_xpdocstring corrected (it claimed a non-existentjax.scipy.special.kvfast-path) and theImportErrormessage now points attfp-nightly. No behaviour change.autogalaxy_workspace_test (1 new file)
scripts/jax_assertions/matern_regularization.py— regression guard: numpy-vs-JAX parity of the Matern regularization matrix (matern_cov_matrix_from+inv_via_cholesky+MaternKernel.regularization_matrix_from), jitted, atnu=0.5andnu=1.7(nu is a continuous free prior, so the Bessel term is tested at non-half-integer order). Lives here, not intest_autoarray/, because library unit tests are NumPy-only by convention.Ship evidence
jax==0.10.2+tensorflow-probability==0.25.0→import tensorflow_probability.substrates.jaxcrashes atops.py:681,AttributeError: module 'jax.interpreters.xla' has no attribute 'pytype_aval_mappings'. Exactly the reported failure.tfp-nightly==0.26.0.dev20260713onjax==0.10.2→bessel_kve(0.5,1.0)=1.2533, import + compute succeed.matern_regularization: all assertions passed.test_matern_kernel.py+test_matern_adapt_kernel.py→ 2 passed (pre-existing divide-by-zero warning inmatern_adapt_kernel.py, unrelated).On your go-ahead
Reply here (or relaunch) and I'll run the full
ship_librarygate (worktreepyteston PyAutoArray + review + Heart), commit, push, and open the library-first PR on PyAutoArray, then theautogalaxy_workspace_testPR behind it. After that, re-run release validation to confirm the Matern failures clear.Trade-off you accepted (AskUserQuestion)
tfp-nightlynightlies are eventually pruned from PyPI, so the pinned build may 404 in the far future. The durable fix — vendoring tfp'sbessel_kveto drop the tfp dependency entirely — is filed as a follow-up (PyAutoMind/draft/refactor/autoarray/matern_vendor_bessel_kve.md, to be created).Jammy2211 commentedon Jul 13, 2026
✅ Shipped (PRs open, pending-release). Human ship sign-off given ("ship all").
pytest test_autoarray/897 passedHeart was RED at ship (all reasons = the release-hold these PRs remediate + stale-evidence/hygiene, not branch breakage); shipped on explicit human go-ahead. Merge/close remain yours. After #386 + PyAutoFit#1366 merge, re-run release validate to confirm GREEN.
fix: pin tfp-nightly for the JAX Matern-kernel path (tfp 0.25.0 × jax…