From f8fc7f3fb12d3a9445b318b93399c0cb3e2ed2bf Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 13 Jul 2026 19:24:56 +0100 Subject: [PATCH] =?UTF-8?q?fix:=20pin=20tfp-nightly=20for=20the=20JAX=20Ma?= =?UTF-8?q?tern-kernel=20path=20(tfp=200.25.0=20=C3=97=20jax=200.10.2=20cr?= =?UTF-8?q?ash)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The last stable tensorflow-probability (0.25.0) references jax.interpreters.xla.pytype_aval_mappings, removed from modern JAX, so it crashes at import under the resolved jax>=0.7 stack — breaking every Matern-kernel pixelization on the JAX backend (release-validation run 29266305445). Swap the pin to tfp-nightly==0.26.0.dev20260713 (verified to compute bessel_kve on jax 0.10.2) and correct the kv_xp docstring/ImportError. Fixes #385 Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TDCsgNCXacXqiWAPTswWCt --- .../inversion/regularization/matern_kernel.py | 16 ++++++++++++---- pyproject.toml | 8 +++++++- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/autoarray/inversion/regularization/matern_kernel.py b/autoarray/inversion/regularization/matern_kernel.py index b71c077b5..fdfed01d8 100644 --- a/autoarray/inversion/regularization/matern_kernel.py +++ b/autoarray/inversion/regularization/matern_kernel.py @@ -17,8 +17,14 @@ def kv_xp(v, z, xp=np): -> scipy.special.kv JAX backend: - -> jax.scipy.special.kv if available - -> else tfp.substrates.jax.math.bessel_kve * exp(-|z|) + -> tensorflow_probability.substrates.jax.math.bessel_kve * exp(-|z|) + + `jax.scipy.special` has no modified-Bessel-K (`kv`/`kve`) of arbitrary real + order, so the JAX path relies on tfp's `bessel_kve`. Note the tfp dependency + is the *nightly* build (`tfp-nightly`): the last stable release + (`tensorflow-probability==0.25.0`) crashes at import under the resolved + `jax>=0.7` stack (it references `jax.interpreters.xla.pytype_aval_mappings`, + removed from modern JAX). """ # ------------------------- @@ -39,8 +45,10 @@ def kv_xp(v, z, xp=np): return tfp.math.bessel_kve(v, z) * xp.exp(-xp.abs(z)) except ImportError: raise ImportError( - "To use the JAX backend with the Matérn kernel, " - "please install tensorflow-probability via `pip install tensorflow-probability==0.25.0`." + "To use the JAX backend with the Matérn kernel, install the " + "tensorflow-probability nightly via " + "`pip install tfp-nightly==0.26.0.dev20260713` " + "(the last stable release, 0.25.0, is incompatible with modern JAX)." ) diff --git a/pyproject.toml b/pyproject.toml index 8053d3e07..bce78d1c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,13 @@ optional = [ "numba", "nufftax>=0.4.0,<0.5.0; python_version >= '3.12'", "pynufft", - "tensorflow-probability==0.25.0" + # tfp provides the modified-Bessel `bessel_kve` used by the JAX Matern-kernel + # regularization path (autoarray/inversion/regularization/matern_kernel.py). + # The last *stable* release (tensorflow-probability==0.25.0) targets a JAX where + # `jax.interpreters.xla.pytype_aval_mappings` still existed; it crashes at import + # under the resolved `jax>=0.7` stack (autoconf[jax]). The nightly restores + # compatibility. Durable fix (vendor bessel_kve, drop tfp) tracked separately. + "tfp-nightly==0.26.0.dev20260713" ] test = ["pytest"] dev = ["pytest", "black", "numba", "nufftax>=0.4.0,<0.5.0; python_version >= '3.12'", "pynufft==2022.2.2"]