fix: Matern-kernel JAX crash — tfp 0.25.0 incompatible with jax 0.10.2

[Jammy2211]

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_mappings is already gone by jax 0.9.2, and autoconf[jax] floors jax>=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-nightly works: verified tfp-nightly==0.26.0.dev20260713 computes bessel_kve(0.5,1.0)=1.2533 on jax 0.10.2, and installs via an explicit == pin without --pre. ← chosen
• (c) drop tfp, compute kv directly — jax.scipy.special has no kv/kve, and nu is a free Uniform(0.5, 5.5) prior (continuous — no closed-form shortcut). The only route is vendoring tfp's bessel_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.dev20260713 in PyAutoArray's [jax] extra (import name is unchanged; matern_kernel.py needs no code edit).
• Fix the kv_xp docstring (it claims a jax.scipy.special.kv fast-path that doesn't exist) and point the ImportError message at tfp-nightly.
• Add a JAX Matern-regularization parity assertion to autogalaxy_workspace_test/scripts/jax_assertions/ (JAX guards don't belong in test_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_kve and 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

1. PyAutoArray/pyproject.toml ([jax] extra, ~line 61): "tensorflow-probability==0.25.0" → "tfp-nightly==0.26.0.dev20260713".
2. PyAutoArray/autoarray/inversion/regularization/matern_kernel.py: correct the kv_xp docstring and ImportError message (no behaviour change).
3. autogalaxy_workspace_test/scripts/jax_assertions/matern_regularization.py (new): build a MaternKernel regularization matrix for a small mesh with xp=np and with xp=jnp under jax.jit; assert numpy/jax parity (exercises the tfp bessel_kve path).
4. Verify in repro venv.
5. 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.

[Jammy2211]

⏸️ Parked at ship sign-off — supervised --auto (needs your go-ahead to open the PR)

Implementation is done and verified; under the supervised bug 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, branch feature/matern-tfp-jax-compat)

PyAutoArray (2 files)

• pyproject.toml — optional extra: 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_xp docstring corrected (it claimed a non-existent jax.scipy.special.kv fast-path) and the ImportError message now points at tfp-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, at nu=0.5 and nu=1.7 (nu is a continuous free prior, so the Bessel term is tested at non-half-integer order). Lives here, not in test_autoarray/, because library unit tests are NumPy-only by convention.

Ship evidence

• Repro (broken): clean env jax==0.10.2 + tensorflow-probability==0.25.0 → import tensorflow_probability.substrates.jax crashes at ops.py:681, AttributeError: module 'jax.interpreters.xla' has no attribute 'pytype_aval_mappings'. Exactly the reported failure.
• Fixed: same clean env with tfp-nightly==0.26.0.dev20260713 on jax==0.10.2 → bessel_kve(0.5,1.0)=1.2533, import + compute succeed.
• Parity: new assertion script → matern_regularization: all assertions passed.
• Unit tests: test_matern_kernel.py + test_matern_adapt_kernel.py → 2 passed (pre-existing divide-by-zero warning in matern_adapt_kernel.py, unrelated).
• API changes: none (public surface unchanged). Scripts changed: one new workspace_test assertion.

On your go-ahead

Reply here (or relaunch) and I'll run the full ship_library gate (worktree pytest on PyAutoArray + review + Heart), commit, push, and open the library-first PR on PyAutoArray, then the autogalaxy_workspace_test PR behind it. After that, re-run release validation to confirm the Matern failures clear.

Trade-off you accepted (AskUserQuestion)

tfp-nightly nightlies are eventually pruned from PyPI, so the pinned build may 404 in the far future. The durable fix — vendoring tfp's bessel_kve to drop the tfp dependency entirely — is filed as a follow-up (PyAutoMind/draft/refactor/autoarray/matern_vendor_bessel_kve.md, to be created).

[Jammy2211]

✅ Shipped (PRs open, pending-release). Human ship sign-off given ("ship all").

• Library: fix: pin tfp-nightly for JAX Matern-kernel path (tfp 0.25.0 × jax 0.10.2 crash) #386 — pytest test_autoarray/ 897 passed
• Regression: test: JAX Matern-regularization parity assertion (guards PyAutoArray#386) autogalaxy_workspace_test#70 (Matern parity, merge behind fix: pin tfp-nightly for JAX Matern-kernel path (tfp 0.25.0 × jax 0.10.2 crash) #386)

Heart 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.