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"]