feat: xp-ify ExponentialKernel + document JAX-gradient support in docstrings

ExponentialKernel's covariance build now threads xp (it previously dropped
the backend at the call site, raising TracerArrayConversionError under
jax.jit/jax.grad) and computes pairwise distances with the same NaN-safe
sqrt(d^2 + 1e-20) dot-product identity as matern_cov_matrix_from —
linalg.norm's derivative is NaN at the zero diagonal, which would poison
every JAX gradient through the kernel. Verified: unit tests pass and JAX
gradients are finite, non-zero and eager/jit-consistent on both the
rectangular and KNN meshes.

Also records the 2026-07 regularization x mesh gradient-sweep findings
where users look them up — the class docstrings: which schemes are
JAX-differentiable on which mesh family (analytic vs scipy neighbors,
split-family compatibility), the Matern/tfp bessel_kve gradient support and
its tfp-nightly requirement, the kernel schemes' explicit-inverse
conditioning caveat, the adaptive-defaults-are-uniform footnote, and the
Delaunay vs KNN mesh gradient split. Behavioural change is limited to
ExponentialKernel's distance form (round-off level; its only unit test
asserts at 1e-4 and passes unchanged).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy

Author: claude

autoarray/inversion/mesh/mesh/delaunay.py

@@ -25,6 +25,12 @@ def __init__(
         according to their barycentric distances, providing a smooth, piecewise-linear
         reconstruction.
 
+        **JAX & gradient support**: the likelihood runs under ``jax.jit``
+        (the triangulation is host-called via ``jax.pure_callback``), but
+        ``jax.grad`` is unavailable — ``pure_callback`` has no JVP rule. The
+        ``KNearestNeighbor`` / ``KNNBarycentric`` subclasses are the
+        gradient-capable members of this mesh family.
+
         Zeroed pixels
         -------------
         The `zeroed_pixels` parameter specifies a number of mesh vertices that are

autoarray/inversion/mesh/mesh/knn.py

@@ -26,6 +26,16 @@ def __init__(
         Neighbour connections may be further restricted using a distance-based criterion,
         and optionally subdivided to improve stability for highly irregular meshes.
 
+        **JAX & gradient support** (2026-07, FD-certified): the kNN
+        interpolation is pure JAX (blocked brute-force ``lax.top_k`` +
+        Wendland weights — no scipy callback), so unlike the parent
+        ``Delaunay`` mesh the full likelihood is differentiable, with
+        gradients flowing through both the traced query points and the
+        traced mesh vertices. Pair it with a split-family regularization
+        (``ConstantSplit`` / ``AdaptSplit``) or a kernel scheme — the
+        neighbor-based schemes (``Constant`` / ``Adapt``) call scipy on the
+        traced mesh grid and cannot differentiate.
+
         Parameters
         ----------
         pixels : int
@@ -92,6 +102,11 @@ class KNNBarycentric(KNearestNeighbor):
     (``k_neighbors``, ``radius_scale``, ``split_neighbor_division``) are
     inherited and still control the regularization-spacing computation, but the
     *interpolation* weights always use k=3 + barycentric and ignore them.
+    Gradients are FD-certified like the parent (2026-07), but this mesh
+    FAILED its science gate as a Delaunay replacement (PyAutoArray#317 —
+    ~5% of vertices are never any query's nearest-3, drifting the
+    log-evidence by ~2%): use it for gradient experiments, not production
+    science.
     """
 
     @property

autoarray/inversion/regularization/adapt.py

@@ -172,6 +172,15 @@ def __init__(
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07 gradient sweep): as for
+        ``Constant`` — JAX-differentiable and FD-certified on the rectangular
+        mesh family (this is the rectangular production scheme), but raises
+        ``TracerArrayConversionError`` on the Delaunay mesh family, whose
+        neighbors come from a direct scipy call on the traced mesh grid (use
+        ``AdaptSplit`` there). Note the defaults
+        ``inner_coefficient == outer_coefficient == 1.0`` make the weighting
+        uniform — numerically identical to ``Constant(coefficient=1.0)``.
+
         Parameters
         ----------
         coefficients

autoarray/inversion/regularization/adapt_split.py

@@ -61,6 +61,13 @@ def __init__(
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07 gradient sweep): as for
+        ``ConstantSplit`` — differentiable and FD-certified on the Delaunay
+        mesh family (e.g. the KNN meshes), structurally incompatible with the
+        rectangular meshes. Note the defaults
+        ``inner_coefficient == outer_coefficient == 1.0`` make the weighting
+        uniform — numerically identical to ``ConstantSplit(coefficient=1.0)``.
+
         Parameters
         ----------
         coefficients

autoarray/inversion/regularization/adapt_split_zeroth.py

@@ -60,6 +60,12 @@ def __init__(
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07 gradient sweep): not currently
+        JAX-differentiable on either gradient-capable mesh family — the split
+        leg is structurally incompatible with the rectangular meshes (shape
+        error) and the zeroth/adapt legs hit numpy-on-traced-array operations
+        on the Delaunay mesh family (``TracerArrayConversionError``).
+
         Parameters
         ----------
         coefficients

autoarray/inversion/regularization/brightness_zeroth.py

@@ -82,6 +82,11 @@ def __init__(
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07 gradient sweep): not yet xp-ported —
+        the pixel-signal thresholding applies numpy operations to traced
+        arrays, raising ``TracerArrayConversionError`` under ``jax.jit`` /
+        ``jax.grad`` on every mesh.
+
         Parameters
         ----------
         coefficient

autoarray/inversion/regularization/constant.py

@@ -86,6 +86,15 @@ def __init__(self, coefficient: float = 1.0):
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07 gradient sweep): on meshes with an
+        analytic neighbor structure (the rectangular family) this scheme is
+        JAX-differentiable and FD-certified. On the Delaunay mesh family
+        (``Delaunay``, ``KNearestNeighbor``, ``KNNBarycentric``) the neighbors
+        come from a direct ``scipy.spatial.Delaunay`` call on the traced
+        source-plane mesh grid, so it raises ``TracerArrayConversionError``
+        under ``jax.jit`` / ``jax.grad`` — use a split-family scheme
+        (e.g. ``ConstantSplit``) there instead.
+
         Parameters
         ----------
         coefficient

autoarray/inversion/regularization/constant_split.py

@@ -35,6 +35,13 @@ def __init__(self, coefficient: float = 1.0):
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07 gradient sweep): differentiable and
+        FD-certified on the Delaunay mesh family (e.g. the KNN meshes), whose
+        interpolators supply pure-xp split mappings. Structurally incompatible
+        with the rectangular meshes, whose interpolator reuses its per-query
+        4-corner mappings for the split path (shape error) — use ``Constant``
+        or ``Adapt`` there.
+
         Parameters
         ----------
         coefficient

autoarray/inversion/regularization/exponential_kernel.py

@@ -20,30 +20,40 @@ def exp_cov_matrix_from(
 
     with a tiny jitter 1e-8 added on the diagonal for numerical stability.
 
+    The pairwise distances use the ``||x||^2 + ||y||^2 - 2 x.y`` identity with a
+    ``sqrt(d^2 + 1e-20)`` floor (mirroring ``matern_cov_matrix_from``) rather than
+    ``linalg.norm`` of an (N, N, 2) difference cube: the norm's derivative is NaN at
+    the zero diagonal, which would poison every JAX gradient through this kernel.
+
     Parameters
     ----------
     scale
         The length‐scale of the exponential kernel.
     pixel_points
         Array of shape (N, 2) giving the (y,x) coordinates of each source‐plane pixel.
+    xp
+        Backend (numpy or jax.numpy).
 
     Returns
     -------
     np.ndarray, shape (N, N)
         The exponential covariance matrix.
     """
-    # pairwise differences: shape (N, N, 2)
-    diff = pixel_points[:, None, :] - pixel_points[None, :, :]
+    pts = xp.asarray(pixel_points)
+
+    # ||x - y||^2 = ||x||^2 + ||y||^2 - 2 x·y
+    x2 = xp.sum(pts * pts, axis=1, keepdims=True)  # (N, 1)
+    dist_sq = x2 + x2.T - 2.0 * (pts @ pts.T)  # (N, N)
+    dist_sq = xp.maximum(dist_sq, 0.0)  # numerical safety
 
-    # Euclidean distances: shape (N, N)
-    d = xp.linalg.norm(diff, axis=-1)
+    d = xp.sqrt(dist_sq + 1e-20)  # (N, N)
 
     # exponential kernel
     cov = xp.exp(-d / scale)
 
     # add a small jitter on the diagonal
-    N = pixel_points.shape[0]
-    cov = cov + xp.eye(N) * 1e-8
+    N = pts.shape[0]
+    cov = cov + xp.eye(N, dtype=cov.dtype) * 1e-8
 
     return cov
 
@@ -64,6 +74,14 @@ def __init__(self, coefficient: float = 1.0, scale: float = 1.0):
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07): xp-threaded and
+        JAX-differentiable end-to-end — the pairwise distances use the
+        NaN-safe ``sqrt(d^2 + 1e-20)`` form (see ``exp_cov_matrix_from``).
+        Caveat: the regularization matrix is an explicit dense inverse of the
+        kernel covariance, whose conditioning on clustered mesh vertices puts
+        a small numerical noise floor on the likelihood (see ``MaternKernel``
+        for the measured detail).
+
         Parameters
         ----------
         coefficient
@@ -113,6 +131,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
         covariance_matrix = exp_cov_matrix_from(
             scale=self.scale,
             pixel_points=linear_obj.source_plane_mesh_grid.array,
+            xp=xp,
         )
 
         return self.coefficient * xp.linalg.inv(covariance_matrix)

autoarray/inversion/regularization/gaussian_kernel.py

@@ -64,6 +64,12 @@ def __init__(self, coefficient: float = 1.0, scale: float = 1.0):
 
         A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
 
+        **JAX & gradient support** (2026-07 gradient sweep): xp-threaded and
+        JAX-differentiable end-to-end. Caveat: the regularization matrix is an
+        explicit dense inverse of the kernel covariance, whose conditioning on
+        clustered mesh vertices puts a small numerical noise floor on the
+        likelihood (see ``MaternKernel`` for the measured detail).
+
         Parameters
         ----------
         coefficient