feat: unlock jax.grad through the Delaunay mesh via frozen integer tables

Wraps the qhull pure_callback input in stop_gradient in _jax_delaunay_tables.
pure_callback has no JVP rule and previously raised under jax.grad; with the
visibility walk running point location, barycentric weights, dual areas and
split points in-graph from the traced arrays, the callback returns only int32
connectivity tables, which are piecewise-constant in the vertex positions —
their true derivative is zero everywhere except the measure-zero triangle-
flip events, where the likelihood itself is discontinuous and no gradient
exists for any method. Freezing them under differentiation therefore yields
the exact almost-everywhere derivative, not an approximation.

Values are bit-identical (the primal is untouched): the jax_likelihood
delaunay regression literal passes unchanged. FD-certified on the production
shape (Hilbert + edge zeroing + AdaptSplit) by the new
autolens_workspace_test scripts/imaging/jax_grad/delaunay.py — 14/14 params
live, lens light at 1e-8..1e-10, mass/shear at 1e-5..2e-3 (FD steps
straddling flip events; documented rtol=1e-2). Mesh docstrings updated:
Delaunay is now gradient-capable; the KNN meshes' remaining edge is batched
throughput (the tables callback is vmap_method=sequential, one host qhull
call per vmap lane).

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

Author: claude

autoarray/inversion/mesh/interpolator/delaunay.py

@@ -118,7 +118,22 @@ def scipy_delaunay_tri_only(points_np):
 
 
 def _jax_delaunay_tables(points):
-    """Run the qhull-only callback with fixed output shapes."""
+    """Run the qhull-only callback with fixed output shapes.
+
+    The callback input is wrapped in ``stop_gradient`` so ``jax.grad`` /
+    ``jax.value_and_grad`` flow through the Delaunay likelihood
+    (``pure_callback`` has no JVP rule and would otherwise raise). This is
+    NOT an approximation: the callback returns only int32 connectivity
+    tables, which are piecewise-constant in the vertex positions — their
+    true derivative is exactly zero everywhere except the measure-zero
+    re-wiring (triangle-flip) events, where the likelihood itself is
+    discontinuous and no gradient exists for any method. Every quantity
+    with a non-zero derivative (point location via the visibility walk,
+    barycentric weights, dual areas, split points) is computed in-graph
+    from the traced ``points``, so the frozen-tables gradient is the exact
+    almost-everywhere derivative. FD-certified 2026-07-26
+    (autolens_workspace_test ``scripts/imaging/jax_grad/delaunay.py``).
+    """
     import jax
     import jax.numpy as jnp
 
@@ -130,7 +145,7 @@ def _jax_delaunay_tables(points):
             jax.ShapeDtypeStruct((2 * N, 3), jnp.int32),
             jax.ShapeDtypeStruct((N,), jnp.int32),
         ),
-        points,
+        jax.lax.stop_gradient(points),
         vmap_method="sequential",
     )
 

autoarray/inversion/mesh/mesh/delaunay.py

@@ -25,11 +25,17 @@ 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.
+        **JAX & gradient support** (2026-07-26, FD-certified): the likelihood
+        runs under ``jax.jit`` and is differentiable — the host-called qhull
+        ``pure_callback`` returns only integer connectivity tables (frozen
+        under differentiation via ``stop_gradient``; their true derivative is
+        zero between re-wiring events), while point location, barycentric
+        weights, dual areas and split points are computed in-graph from the
+        traced vertices, so ``jax.grad`` returns the exact almost-everywhere
+        derivative. Caveat for batched samplers: the callback is
+        ``vmap_method="sequential"`` (one host qhull call per vmap lane) —
+        the ``KNearestNeighbor`` / ``KNNBarycentric`` subclasses avoid the
+        callback entirely and remain the batched-throughput option.
 
         Zeroed pixels
         -------------

autoarray/inversion/mesh/mesh/knn.py

@@ -28,13 +28,15 @@ def __init__(
 
         **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.
+        Wendland weights — no scipy callback), so the full likelihood is
+        differentiable, with gradients flowing through both the traced query
+        points and the traced mesh vertices. (The parent ``Delaunay`` mesh
+        is also differentiable as of 2026-07-26 via its frozen integer
+        tables; this mesh's remaining edge is batched throughput — no
+        per-vmap-lane host callback.) 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
         ----------