diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 77911fd..80f18a3 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/docs/diagonal-masking.md b/docs/diagonal-masking.md new file mode 100644 index 0000000..dcabba2 --- /dev/null +++ b/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. diff --git a/powersig/jax/algorithm.py b/powersig/jax/algorithm.py index d1c7201..5e0a949 100644 --- a/powersig/jax/algorithm.py +++ b/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)*(d-cols +1) - s_start = (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)*(diagonal_index-cols +1) - s_start = (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 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)) diff --git a/powersig/jax/diagonal.py b/powersig/jax/diagonal.py new file mode 100644 index 0000000..bb3b13e --- /dev/null +++ b/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 diff --git a/tests/test_diagonal_masking.py b/tests/test_diagonal_masking.py new file mode 100644 index 0000000..e0a09ce --- /dev/null +++ b/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)))