fix: pin tfp-nightly for the JAX Matern-kernel path (tfp 0.25.0 × jax 0.10.2 crash) (#386)

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


Claude-Session: https://claude.ai/code/session_01TDCsgNCXacXqiWAPTswWCt

Co-authored-by: Jammy2211 <JNightingale2211@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: 3 people

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)."
             )
 
 

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