Fix inactive JAX diagonal cells poisoning path gradients

.github/workflows/ci.yml

@@ -29,5 +29,6 @@ jobs:
       - name: Run tests
         env:
           JAX_PLATFORMS: cpu
+          NVIDIA_TF32_OVERRIDE: "0"
         run: |
-          pytest tests/test_core_jax.py -v
+          pytest tests/test_core_jax.py tests/test_diagonal_masking.py -v

docs/diagonal-masking.md

@@ -0,0 +1,38 @@
+# Inactive diagonal cells and numerical safety
+
+A fixed-width diagonal sweep contains lanes outside the rectangular path-pair
+grid, especially at the first and last diagonals. These lanes are not kernel
+tiles. Evaluating them with wrapped/clamped path indices can repeatedly reuse
+a large increment, overflowing internal polynomial states even when every
+real tile and the final kernel value are finite.
+
+Mask the inputs **before** the static kernel and power-series recurrence:
+use in-bounds gather indices, select boundary conditions with `where` rather
+than multiplication by a Boolean, and give inactive lanes zero boundary
+vectors and zero increments. Masking only the final output is insufficient
+for differentiation: an unused infinity can still produce `0 * Inf = NaN`
+in a pullback. Gradient clipping after backpropagation cannot recover those
+NaNs.
+
+The custom-VJP implementation, when present, must use the same rule in its
+checkpointed forward sweep, forward replay, local adjoints, and path-gradient
+scatter. Inactive lanes must never contribute to either boundary or path
+adjoints. This is geometric masking, not gradient clipping or sanitizing
+non-finite values in real tiles. Genuine overflow in an active tile remains
+visible.
+
+## Reproduction and verification
+
+The tests use only synthetic FP32 paths: a 64-knot scalar query from zero to
+one, and a bank path with one jump from zero to 256 followed by constant
+knots, at numerical order 9. The unpatched sweep gives a finite kernel value
+but non-finite path gradients; the masked implementation agrees with an
+independent row-major scan that visits only real tiles.
+
+The reference checks the same finite-order numerical recurrence, not the
+accuracy of that recurrence against the exact signature kernel for large
+increments. Cancellation in the stress-case gradient permits eight FP32 ULPs
+at the largest component's scale; ordinary small-path tests retain strict
+elementwise tolerances. Additional tests cover linear and RBF kernels,
+unequal lengths in both orientations, one-increment paths, chunked and
+unchunked sweeps, and genuinely overflowing active tiles.

powersig/jax/algorithm.py

@@ -16,6 +16,7 @@
 from tqdm.auto import tqdm
 
 from powersig.jax.jax_series import jax_compute_derivative, jax_compute_derivative_batch
+from powersig.jax.diagonal import diagonal_tile_inputs
 
 
 class PowerSigJax:
@@ -74,8 +75,8 @@ def compute_signature_kernel(self, X: jnp.ndarray, Y: jnp.ndarray, device=None)
         """
         # dX = jax_compute_derivative(X.squeeze(0))
         # dY = jax_compute_derivative(Y.squeeze(0))
-        # Ensure exponents are on the same device as input
-        self.exponents = jax.device_put(self.exponents, device)
+        # Do not mutate instance constants inside jit: that leaks a tracer
+        # when this instance is subsequently reused for another shape or VJP.
          # Calculate values we need before padding
         diagonal_count = ( X.shape[0] -1) + (Y.shape[0] - 1) - 1
         longest_diagonal = min(X.shape[0] - 1, Y.shape[0] - 1)
@@ -277,9 +278,7 @@ def compute_gram_entry(
         def compute_diagonal(d, carry):
             S_buf, T_buf = carry
             # s_start, t_start, dlen = get_diagonal_range(d, dX_i.shape[0], dY_j.shape[0])
-            t_start = (d<cols)*0 + (d>=cols)*(d-cols +1)
-            s_start = (d<cols)*d + (d>=cols)*(cols - 1)
-            dlen = jnp.minimum(rows - t_start, s_start + 1)
+            s_start, t_start, dlen = get_diagonal_range(d, rows, cols)
             is_before_wrap = d < rows
             # dX_L = dX_i.shape[0] - (s_start + 1)
 
@@ -292,16 +291,9 @@ def compute_diagonal(d, carry):
             #                      precision=jax.lax.Precision.HIGHEST)
 
             def next_diagonal_entry(diagonal_index, S, T):
-                # Combine the first two where statements into a single mask
-                s_index = diagonal_index - is_before_wrap
-                t_index = diagonal_index + (1 - is_before_wrap)
-
-                # Avoid branching
-                s = ((t_start + diagonal_index == 0) * self.ic) + ((t_start + diagonal_index != 0) * S[s_index])
-                t = ((s_start - diagonal_index == 0) * self.ic) + ((s_start - diagonal_index != 0) * T[t_index])
-                dX_idx = (s_start - diagonal_index) * ((s_start - diagonal_index) < rows)
-                dY_idx = (t_start + diagonal_index) * ((t_start + diagonal_index) < cols)
-                rho = self.static_kernel(X_i[dX_idx+1],X_i[dX_idx], Y_j[dY_idx+1],Y_j[dY_idx])
+                s, t, rho, _, _ = diagonal_tile_inputs(
+                    diagonal_index, s_start, t_start, dlen, is_before_wrap,
+                    X_i, Y_j, S, T, self.ic, self.static_kernel)
                 # rho = (X_i.shape[0]-1)*(Y_j.shape[0]-1)*jnp.dot(X_i[dX_idx+1]-X_i[dX_idx], Y_j[dY_idx+1]-Y_j[dY_idx], precision = jax.lax.Precision.HIGHEST)
                 # rho = jnp.dot(dX_i[dX_idx], dY_j[dY_idx], precision = jax.lax.Precision.HIGHEST)
                 # jax.debug.print("""
@@ -410,24 +402,16 @@ def chunked_compute_gram_entry(
             # print(f"batch_longest_diag = {batch_longest_diag}")
             def next_diagonal(diagonal_index,carry):
                 # jax.debug.print("========================= START OF BATCH {} =========================\n", d)
-                t_start = (diagonal_index<cols)*0 + (diagonal_index>=cols)*(diagonal_index-cols +1)
-                s_start = (diagonal_index<cols)*diagonal_index + (diagonal_index>=cols)*(cols - 1)
+                s_start, t_start, dlen = get_diagonal_range(diagonal_index, rows, cols)
 
                 is_before_wrap = diagonal_index < rows
                 # rho = jax_compute_dot_prod_batch(jnp.take(dX_i, s_start-diagonal_indices, axis=0, fill_value=0), jnp.take(dY_j, t_start+diagonal_indices, axis=0, fill_value=0))
                 # rho = jnp.einsum('ij,ij->i', jnp.take(dX_i, s_start-diagonal_indices, axis=0, fill_value=0), jnp.take(dY_j, t_start+diagonal_indices, axis=0, fill_value=0),
                 #                  precision=jax.lax.Precision.HIGHEST)
                 def next_diagonal_entry(index_in_diagonal, S, T):
-                    # Combine the first two where statements into a single mask
-                    s_index = index_in_diagonal - is_before_wrap
-                    t_index = index_in_diagonal + (1 - is_before_wrap)
-
-                    # Avoid branching
-                    s = ((t_start + index_in_diagonal == 0) * self.ic) + ((t_start + index_in_diagonal != 0) * S[s_index])
-                    t = ((s_start - index_in_diagonal == 0) * self.ic) + ((s_start - index_in_diagonal != 0) * T[t_index])
-                    dX_idx = (s_start - index_in_diagonal) * ((s_start - index_in_diagonal) < rows)
-                    dY_idx = (t_start + index_in_diagonal) * ((t_start + index_in_diagonal) < cols)
-                    rho = self.static_kernel(X_i[dX_idx+1],X_i[dX_idx], Y_j[dY_idx+1],Y_j[dY_idx])
+                    s, t, rho, _, _ = diagonal_tile_inputs(
+                        index_in_diagonal, s_start, t_start, dlen, is_before_wrap,
+                        X_i, Y_j, S, T, self.ic, self.static_kernel)
                     # rho = (X_i.shape[0]-1)*(Y_j.shape[0]-1)*jnp.dot(X_i[dX_idx+1]-X_i[dX_idx], Y_j[dY_idx+1]-Y_j[dY_idx], precision = jax.lax.Precision.HIGHEST)
                     # rho = jnp.dot(dX_i[dX_idx], dY_j[dY_idx], precision = jax.lax.Precision.HIGHEST)
                     # jax.debug.print("""
@@ -810,18 +794,9 @@ def process_column(c):
 @jit
 def get_diagonal_range(d: int, rows: int, cols: int) -> Tuple[int, int, int]:
     # d, s_start, t_start are 0 based indexes while rows/cols are shapes.
-    t_start = jnp.where(d<cols, 0, d-cols +1)
-    s_start = jnp.where(d<cols, d, cols - 1)
-    dlen = jnp.minimum(rows - t_start, s_start + 1)
-    # if d < cols:
-    #     # if d < cols, then we haven't hit the right edge of the grid
-    #     t_start = 0
-    #     s_start = d
-    # else:
-    #     # if d >= cols then we have the right edge and wrapped around the corner
-    #     t_start = d - cols + 1  # diag index - cols + 1
-    #     s_start = cols - 1
-    # return s_start, t_start, min(rows - t_start, s_start + 1)
+    s_start = jnp.minimum(d, rows - 1)
+    t_start = d - s_start
+    dlen = jnp.minimum(s_start + 1, cols - t_start)
     return s_start, t_start, dlen
 
 # @partial(jit, static_argnums=(1,2,3))

powersig/jax/diagonal.py

@@ -0,0 +1,33 @@
+"""Safe inputs for fixed-width JAX diagonal sweeps.
+
+Inactive SIMD lanes are not signature-kernel tiles. Mask before arithmetic,
+not just at the output: an unused overflowing state can otherwise produce
+0 * Inf in a reverse pass and contaminate real path gradients.
+"""
+import jax.numpy as jnp
+
+
+def masked_path_points(X, Y, x_index, y_index, active):
+    """Finite dummy inputs for inactive lanes, without changing active tiles."""
+    return (
+        jnp.where(active, X[x_index + 1], 0),
+        jnp.where(active, X[x_index], 0),
+        jnp.where(active, Y[y_index + 1], 0),
+        jnp.where(active, Y[y_index], 0),
+    )
+
+
+def diagonal_tile_inputs(index, s_start, t_start, length, before_wrap,
+                         X, Y, S, T, ic, static_kernel):
+    """Gather one tile with in-bounds indices and zeroed inactive operands."""
+    active = index < length
+    s_index = jnp.clip(index - before_wrap, 0, S.shape[0] - 1)
+    t_index = jnp.clip(index + (1 - before_wrap), 0, T.shape[0] - 1)
+    s = jnp.where(t_start + index == 0, ic, S[s_index])
+    t = jnp.where(s_start - index == 0, ic, T[t_index])
+    s, t = jnp.where(active, s, 0), jnp.where(active, t, 0)
+    x_index = jnp.clip(s_start - index, 0, X.shape[0] - 2)
+    y_index = jnp.clip(t_start + index, 0, Y.shape[0] - 2)
+    rho = static_kernel(*masked_path_points(X, Y, x_index, y_index, active))
+    rho = jnp.where(active, rho, 0)
+    return s, t, rho, x_index, y_index

tests/test_diagonal_masking.py

@@ -0,0 +1,121 @@
+"""Dataset-free FP32 regressions for inactive diagonal lanes."""
+import jax
+import jax.numpy as jnp
+from jax.scipy.linalg import toeplitz
+import numpy as np
+import pytest
+
+from powersig.jax.algorithm import PowerSigJax
+from powersig.jax.diagonal import diagonal_tile_inputs
+from powersig.jax.static_kernels import linear_kernel, rbf_kernel
+
+
+@pytest.fixture(autouse=True)
+def strict_float32():
+    old_x64 = jax.config.jax_enable_x64
+    old_precision = jax.config.jax_default_matmul_precision
+    jax.config.update("jax_enable_x64", False)
+    jax.config.update("jax_default_matmul_precision", "highest")
+    yield
+    jax.config.update("jax_enable_x64", old_x64)
+    jax.config.update("jax_default_matmul_precision", old_precision)
+
+
+def make_ps(kernel=linear_kernel):
+    return PowerSigJax(order=9, dtype=jnp.float32,
+                       device=jax.devices()[0], static_kernel=kernel)
+
+
+def overflow_paths():
+    # A single large bank increment, followed by constant knots. All REAL
+    # tile states and derivatives are finite. The old fixed-width sweep
+    # repeatedly reuses this increment in nonexistent tiles and overflows.
+    x = jnp.linspace(0, 1, 64, dtype=jnp.float32)[:, None]
+    y = jnp.full_like(x, 256).at[0].set(0)
+    return x, y
+
+
+def dense_reference(ps, x, y):
+    """Row-major scan over only real tiles; no diagonal padding or masking."""
+    cols = len(y) - 1
+    bottom = jnp.tile(ps.ic, (cols, 1))
+
+    def row(bottom_edges, i):
+        def tile(left_edge, args):
+            j, bottom_edge = args
+            rho = ps.static_kernel(x[i + 1], x[i], y[j + 1], y[j])
+            r = rho ** ps.exponents
+            matrix = ps.psi_t * toeplitz(bottom_edge, left_edge)
+            right = r @ jnp.triu(matrix, 1) + (ps.psi_s @ bottom_edge) * r
+            top = (ps.psi_s @ left_edge) * r + jnp.tril(matrix, -1) @ r
+            return right, top
+        right, top_edges = jax.lax.scan(tile, ps.ic,
+                                       (jnp.arange(cols), bottom_edges))
+        return top_edges, right
+
+    _, rights = jax.lax.scan(row, bottom, jnp.arange(len(x) - 1))
+    return jnp.sum(rights[-1])
+
+
+def assert_close(actual, expected, cancellation=False):
+    for got, want in zip(jax.tree.leaves(actual), jax.tree.leaves(expected)):
+        assert got.dtype == jnp.float32
+        assert bool(jnp.isfinite(got).all())
+        assert bool(jnp.isfinite(want).all())
+        # The stress case subtracts O(1e11) terms into near-zero coordinates.
+        # Allow eight ULPs of the largest component ONLY for that case; keep
+        # strict elementwise tolerances for ordinary linear/RBF paths.
+        atol = 2e-5
+        if cancellation:
+            atol = max(atol, 8 * float(np.spacing(np.max(np.abs(np.asarray(want))))))
+        np.testing.assert_allclose(got, want, rtol=3e-4, atol=atol)
+
+
+@pytest.mark.parametrize("chunked", [False, True])
+def test_unused_lanes_cannot_overflow_value_or_native_gradient(chunked):
+    ps = make_ps()
+    x, y = overflow_paths()
+    forward = (ps.compute_signature_kernel_chunked if chunked
+               else ps.compute_signature_kernel)
+    got = jax.jit(jax.value_and_grad(forward, argnums=(0, 1)))(x, y)
+    want = jax.jit(jax.value_and_grad(lambda a, b: dense_reference(ps, a, b),
+                                    argnums=(0, 1)))(x, y)
+    assert_close(got, want, cancellation=True)
+
+
+@pytest.mark.parametrize("lengths", [(2, 2), (2, 7), (7, 2), (5, 9), (9, 5)])
+@pytest.mark.parametrize("kernel", [linear_kernel, rbf_kernel])
+@pytest.mark.parametrize("chunked", [False, True])
+def test_rectangular_values_and_gradients(lengths, kernel, chunked):
+    ps = make_ps(kernel)
+    rng = np.random.default_rng(123)
+    x = jnp.asarray(rng.normal(size=(lengths[0], 2)).astype(np.float32) * np.float32(.15))
+    y = jnp.asarray(rng.normal(size=(lengths[1], 2)).astype(np.float32) * np.float32(.15))
+    forward = (ps.compute_signature_kernel_chunked if chunked
+               else ps.compute_signature_kernel)
+    got = jax.jit(jax.value_and_grad(forward, argnums=(0, 1)))(x, y)
+    want = jax.jit(jax.value_and_grad(lambda a, b: dense_reference(ps, a, b),
+                                    argnums=(0, 1)))(x, y)
+    assert_close(got, want)
+    np.testing.assert_allclose(forward(x, y), forward(y, x), rtol=3e-5, atol=2e-6)
+
+
+@pytest.mark.parametrize("kernel", [linear_kernel, rbf_kernel])
+def test_poisoned_unused_inputs_are_masked_before_arithmetic(kernel):
+    ps = make_ps(kernel)
+    x = jnp.full((4, 2), jnp.float32(1e30))
+    s = jnp.full((3, 9), jnp.inf, dtype=jnp.float32)
+    t = jnp.full((3, 9), jnp.nan, dtype=jnp.float32)
+    left, bottom, rho, xi, yi = diagonal_tile_inputs(
+        jnp.int32(2), jnp.int32(0), jnp.int32(0), jnp.int32(1), True,
+        x, -x, s, t, ps.ic, ps.static_kernel)
+    np.testing.assert_array_equal(left, jnp.zeros_like(left))
+    np.testing.assert_array_equal(bottom, jnp.zeros_like(bottom))
+    assert float(rho) == 0
+    assert 0 <= int(xi) < 3 and 0 <= int(yi) < 3
+
+
+def test_mask_does_not_hide_overflow_in_a_real_tile():
+    ps = make_ps()
+    x = jnp.array([[0.], [1e10]], dtype=jnp.float32)
+    assert not bool(jnp.isfinite(ps.compute_signature_kernel(x, x)))