From a1ae5d6cdec34ab661b8fd3317ef8a40c235ed05 Mon Sep 17 00:00:00 2001 From: Matthew Tamayo-Rios Date: Wed, 26 Aug 2026 03:06:32 -0700 Subject: [PATCH 1/5] Fix asymmetric diagonal sweep geometry and a jit tracer leak in the JAX backend The diagonal sweep keyed s_start/t_start off `cols`, which walks s past the bottom edge of the grid whenever rows != cols. Anti-diagonal d covers the cells {(s, t) : s + t == d}; s_start is the largest such s and t_start the smallest such t, so s must pin at rows - 1, not cols - 1. Checked against a brute-force enumeration of the grid, the old convention was wrong for 110 of the (d, rows, cols) triples with rows, cols <= 6, and every one of them had rows != cols. The effect was a silently wrong signature kernel for any pair of paths sampled at different lengths. The kernel depends on the paths, not on how finely they are sampled, so upsampling a piecewise-linear path along its own segments must leave the value unchanged; for a 6-point and a 5-point path regridded to a common 21 points, the kernel was off by 2.7e-2 and the error did not shrink with truncation order. Separately, compute_signature_kernel is jitted but assigned to self.exponents, which stores a tracer on the instance and leaks it into every later call -- a compute_signature_kernel followed by a compute_gram_matrix on the same instance died with InvalidInputException. self.exponents is already placed in __init__. Both fixes come with regression tests; all five fail on the previous code. --- powersig/jax/algorithm.py | 28 +++++++----- powersig/util/grid.py | 15 +++---- tests/test_core_jax.py | 94 +++++++++++++++++++++++++++++++++++++-- 3 files changed, 115 insertions(+), 22 deletions(-) diff --git a/powersig/jax/algorithm.py b/powersig/jax/algorithm.py index d1c7201..9393b69 100644 --- a/powersig/jax/algorithm.py +++ b/powersig/jax/algorithm.py @@ -74,8 +74,10 @@ 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) + # NB: do not device_put onto `self` here -- this method is jitted, so + # assigning to an attribute stores a tracer on the instance and leaks it + # into every later call (a subsequent compute_gram_matrix then dies with + # InvalidInputException). `self.exponents` is already placed in __init__. # 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) @@ -276,10 +278,14 @@ 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) + # Anti-diagonal d covers the cells {(s, t) : s + t == d} inside the + # rows x cols grid. s_start is the largest such s and t_start the + # smallest such t, so the sweep pins s at the bottom edge (rows - 1) + # once it runs off it -- keying this off `cols` walks s out of bounds + # whenever rows != cols. See tests/test_core_jax.py::TestDiagonalRange. + s_start = (d=rows)*(rows - 1) + t_start = (d=rows)*(d-rows + 1) + dlen = jnp.minimum(s_start + 1, cols - t_start) is_before_wrap = d < rows # dX_L = dX_i.shape[0] - (s_start + 1) @@ -410,8 +416,8 @@ 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 = (diagonal_index=rows)*(rows - 1) + t_start = (diagonal_index=rows)*(diagonal_index-rows + 1) 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)) @@ -810,9 +816,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 Tuple[int, int, int]: # d, s_start, t_start are 0 based indexes while rows/cols are shapes. - - if d < cols: - # if d < cols, then we haven't hit the right edge of the grid - t_start = 0 + if d < rows: + # We have not yet hit the bottom edge of the grid. s_start = d + t_start = 0 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 + # Once we reach the bottom edge, keep s pinned and advance t. + s_start = rows - 1 + t_start = d - rows + 1 - return s_start, t_start, min(rows - t_start, s_start + 1) + return s_start, t_start, min(s_start + 1, cols - t_start) diff --git a/tests/test_core_jax.py b/tests/test_core_jax.py index 4b5d7d7..1209f21 100644 --- a/tests/test_core_jax.py +++ b/tests/test_core_jax.py @@ -93,15 +93,49 @@ def test_square_grid(self): self.assertEqual((s, t, dlen), (2, 0, 3)) def test_rectangular_grid(self): - # 2 rows, 4 cols + # 2 rows, 4 cols. s_start is the largest s on the anti-diagonal and + # t_start the smallest t, so both stay inside the grid once the sweep + # runs off the bottom edge. s, t, dlen = get_diagonal_range(0, 2, 4) self.assertEqual((s, t, dlen), (0, 0, 1)) s, t, dlen = get_diagonal_range(3, 2, 4) - self.assertEqual((s, t, dlen), (3, 0, 2)) + self.assertEqual((s, t, dlen), (1, 2, 2)) s, t, dlen = get_diagonal_range(4, 2, 4) - self.assertEqual((s, t, dlen), (3, 1, 1)) + self.assertEqual((s, t, dlen), (1, 3, 1)) + + def test_tall_rectangular_grid(self): + # 3 rows, 2 cols -- the transpose of the wide case above. + expected = [ + (0, 0, 1), + (1, 0, 2), + (2, 0, 2), + (2, 1, 1), + ] + self.assertEqual([get_diagonal_range(d, 3, 2) for d in range(4)], expected) + + def test_matches_brute_force_geometry(self): + # Ground truth: enumerate the cells on each anti-diagonal directly. + for rows in range(1, 7): + for cols in range(1, 7): + for d in range(rows + cols - 1): + cells = [ + (s, t) + for s in range(rows) + for t in range(cols) + if s + t == d + ] + expected = ( + max(s for s, _ in cells), + min(t for _, t in cells), + len(cells), + ) + self.assertEqual( + get_diagonal_range(d, rows, cols), + expected, + msg=f"d={d} rows={rows} cols={cols}", + ) # --------------------------------------------------------------------------- @@ -222,5 +256,59 @@ def test_2x2(self): self.assertFalse(jnp.allclose(result[:2], jnp.zeros_like(result[:2]))) +# --------------------------------------------------------------------------- +# Instance reuse across the jitted entry points +# --------------------------------------------------------------------------- +class TestInstanceReuse(unittest.TestCase): + def test_signature_kernel_then_gram_matrix(self): + """compute_signature_kernel is jitted; it must not leave a tracer on self. + + Regression test: assigning to self.exponents inside the jitted method + used to poison the instance, so a later compute_gram_matrix raised + InvalidInputException on a leaked JitTracer. + """ + ps = PowerSigJax(order=8) + path = jnp.asarray(np.linspace(0.0, 1.0, 17).reshape(17, 1)) + first = float(ps.compute_signature_kernel(path, path)) + + batch = jnp.asarray(np.linspace(0.0, 1.0, 17).reshape(1, 17, 1)) + gram = np.asarray(ps.compute_gram_matrix(batch, batch)) + + self.assertEqual(gram.shape, (1, 1)) + np.testing.assert_allclose(gram[0, 0], first, rtol=1e-10, atol=1e-12) + + def test_asymmetric_path_lengths_match_equal_length_reference(self): + """The kernel depends on the paths, not on how finely they are sampled. + + Upsampling a piecewise-linear path along its own segments leaves the path + unchanged, so a 6-vs-5 point pair must agree with the same two paths + re-gridded to a common length. This fails if the diagonal sweep keys its + geometry off `cols` when rows != cols. + """ + + def upsample(P, factor): + out = [] + for i in range(len(P) - 1): + for k in range(factor): + out.append(P[i] + (P[i + 1] - P[i]) * k / factor) + out.append(P[-1]) + return np.array(out) + + rng = np.random.default_rng(123) + X = 0.2 * rng.normal(size=(6, 2)) # 5 segments + Y = 0.2 * rng.normal(size=(5, 2)) # 4 segments + # lcm(5, 4) = 20 -> both become 21 points tracing the identical paths. + Xr, Yr = upsample(X, 4), upsample(Y, 5) + self.assertEqual(Xr.shape, Yr.shape) + + asymmetric = float( + PowerSigJax(order=16).compute_signature_kernel(jnp.asarray(X), jnp.asarray(Y)) + ) + reference = float( + PowerSigJax(order=16).compute_signature_kernel(jnp.asarray(Xr), jnp.asarray(Yr)) + ) + np.testing.assert_allclose(asymmetric, reference, rtol=1e-10, atol=1e-12) + + if __name__ == "__main__": unittest.main() From 2de91d9f61ba97b69442b6db55187415152ea1a5 Mon Sep 17 00:00:00 2001 From: Matthew Tamayo-Rios Date: Wed, 26 Aug 2026 03:06:44 -0700 Subject: [PATCH 2/5] Replace the PyTorch backend with a working implementation and cover it in CI The PyTorch backend on main had never run. PowerSigTorch had no __call__; compute_signature_kernel passed a dtype into build_stencil_s's device parameter and died in dynamo on torch.ones(..., device=torch.float64); compute_gram_matrix was a copy-paste of the JAX version, calling jnp.zeros (never imported) and .at[0].set(1) on a torch tensor; and compute_gram_entry returned from inside its outer loop after the first chunk, with debug prints still in place. None of this was caught because there were no PyTorch tests and CI ran only tests/test_core_jax.py. This takes the implementation from feature/custom-autodiff, which reaches API parity with the JAX backend -- __call__, compute_gram_matrix, pluggable static kernels, device/dtype selection, block sizing, and autodiff -- along with its three test modules. Verified against the JAX backend and against closed form: a linear path matches I_0(2) to 1e-12, Gram matrices agree with JAX to 5e-15 across linear and RBF static kernels, CUDA agrees with CPU to 4e-15, and reverse-mode gradients match central differences to 1.4e-9 relative. CI now runs a backend matrix so neither implementation can rot silently again. The torch job installs the JAX CPU wheel too, since the PyTorch tests cross-check their results against the JAX reference. --- .github/workflows/ci.yml | 26 +- powersig/torch/__init__.py | 74 +- powersig/torch/algorithm.py | 1243 ++++++++++++----------------- powersig/torch/autodiff.py | 503 ++++++++++++ powersig/torch/static_kernels.py | 26 + tests/test_autodiff_torch.py | 126 +++ tests/test_core_torch.py | 286 +++++++ tests/test_prefix_family_torch.py | 184 +++++ 8 files changed, 1705 insertions(+), 763 deletions(-) create mode 100644 powersig/torch/autodiff.py create mode 100644 powersig/torch/static_kernels.py create mode 100644 tests/test_autodiff_torch.py create mode 100644 tests/test_core_torch.py create mode 100644 tests/test_prefix_family_torch.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 77911fd..1b8b508 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,8 +10,12 @@ jobs: test: runs-on: ubuntu-latest strategy: + fail-fast: false matrix: python-version: ["3.12"] + backend: [jax, torch] + + name: test (${{ matrix.backend }}, py${{ matrix.python-version }}) steps: - uses: actions/checkout@v4 @@ -21,13 +25,29 @@ jobs: with: python-version: ${{ matrix.python-version }} + # The torch job installs the JAX CPU wheel too: the torch backend's tests + # cross-check their results against the JAX reference implementation. - name: Install dependencies run: | python -m pip install --upgrade pip - pip install ".[jax-cpu,dev]" + if [ "${{ matrix.backend }}" = "torch" ]; then + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install ".[jax-cpu,dev]" + else + pip install ".[jax-cpu,dev]" + fi + + - name: Run JAX tests + if: matrix.backend == 'jax' + env: + JAX_PLATFORMS: cpu + run: pytest tests/test_core_jax.py -v - - name: Run tests + - name: Run PyTorch tests + if: matrix.backend == 'torch' env: JAX_PLATFORMS: cpu run: | - pytest tests/test_core_jax.py -v + pytest tests/test_core_torch.py \ + tests/test_autodiff_torch.py \ + tests/test_prefix_family_torch.py -v diff --git a/powersig/torch/__init__.py b/powersig/torch/__init__.py index a635531..ee8cd1d 100644 --- a/powersig/torch/__init__.py +++ b/powersig/torch/__init__.py @@ -1,27 +1,65 @@ -""" -PowerSig Torch Module +"""PyTorch backend for signature-kernel computations.""" -This module provides PyTorch implementations of signature kernels and related utilities. -""" - -from .algorithm import PowerSigTorch +from .algorithm import ( + DIAGONAL_CHUNK_SIZE, + JIT_BOUNDARY_THRESHOLD, + PowerSigTorch, + _round_to_power_of_2, + build_psi_stencil, + build_stencil, + build_stencil_s, + build_stencil_t, + compute_block_size, + compute_vandermonde_vectors, + estimate_bytes_per_pair, + get_available_gpu_memory, + get_diagonal_range_bool, + get_max_block_size, +) +from .autodiff import ( + compute_gram_fast_diff, + compute_prefix_family, + compute_prefix_family_fast_diff, + compute_sig_kernel_fast_diff, +) +from .static_kernels import linear_kernel, rbf_fn, rbf_kernel from .utils import ( + chebychev_clamp_and_map, + chebychev_transformation, fractional_brownian_motion, + scale_and_shift, torch_compute_differences, - unity_transform, unity_clamp_and_map, - scale_and_shift, - chebychev_transformation, - chebychev_clamp_and_map + unity_transform, ) __all__ = [ - 'PowerSigTorch', - 'fractional_brownian_motion', - 'torch_compute_differences', - 'unity_transform', - 'unity_clamp_and_map', - 'scale_and_shift', - 'chebychev_transformation', - 'chebychev_clamp_and_map' + "DIAGONAL_CHUNK_SIZE", + "JIT_BOUNDARY_THRESHOLD", + "PowerSigTorch", + "_round_to_power_of_2", + "build_psi_stencil", + "build_stencil", + "build_stencil_s", + "build_stencil_t", + "chebychev_clamp_and_map", + "chebychev_transformation", + "compute_block_size", + "compute_gram_fast_diff", + "compute_prefix_family", + "compute_prefix_family_fast_diff", + "compute_sig_kernel_fast_diff", + "compute_vandermonde_vectors", + "estimate_bytes_per_pair", + "fractional_brownian_motion", + "get_available_gpu_memory", + "get_diagonal_range_bool", + "get_max_block_size", + "linear_kernel", + "rbf_fn", + "rbf_kernel", + "scale_and_shift", + "torch_compute_differences", + "unity_clamp_and_map", + "unity_transform", ] diff --git a/powersig/torch/algorithm.py b/powersig/torch/algorithm.py index c661355..fcc8d61 100644 --- a/powersig/torch/algorithm.py +++ b/powersig/torch/algorithm.py @@ -1,784 +1,543 @@ -from math import ceil, sqrt -from typing import Optional, Tuple +from typing import Callable, Optional, Tuple + import torch -import torch._dynamo +import torch.nn.functional as F +from tqdm.auto import tqdm from powersig.util.grid import get_diagonal_range -from powersig.torch.series import torch_compute_derivative, torch_compute_derivative_batch, torch_compute_dot_prod_batch +from . import static_kernels -torch._dynamo.config.capture_scalar_outputs = True -DIAGONAL_CHUNK_SIZE = 16 -class PowerSigTorch: - def __init__(self, order: int = 32, device: Optional[torch.device] = None): - # Select device - prefer CUDA if available, otherwise use CPU - self.order = order - if device is None: - devices = torch.cuda.device_count() - self.device = torch.device("cuda:1" if devices == 2 else "cuda" if devices >0 else "cpu") - else: - self.device = device - # self.exponents = jnp.arange(self.order) - self.exponents = build_increasing_matrix(self.order, dtype=torch.int8, device=self.device) - - @torch.compile(mode="max-autotune", fullgraph=True) - def compute_signature_kernel(self, X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor: - """ - Compute the signature kernel between two sets of time series. - Args: - X: torch tensor of shape (length, dim) representing the first set of time series - Y: torch tensor of shape (length, dim) representing the second set of time series - symmetric: If True, computes the kernel matrix for the combined set of X and Y. Default is False. - - Returns: - A float representing the signature kernel between X and Y - - """ - dX = torch_compute_derivative(X.squeeze(0)) - dY = torch_compute_derivative(Y.squeeze(0)) - # Calculate values we need before padding - diagonal_count = dX.shape[0] + dY.shape[0] - 1 - longest_diagonal = min(dX.shape[0], dY.shape[0]) - indices = torch.arange(longest_diagonal) - ic = torch.zeros([ self.order], dtype=dX.dtype) - ic[0] = 1 - diagonal_batch_size = ceil(sqrt(longest_diagonal)) - # Generate Vandermonde vectors with high precision - ds = 1.0 / dX.shape[0] - dt = 1.0 / dY.shape[0] - v_s, v_t = compute_vandermonde_vectors(ds, dt, self.order, dX.dtype, dX.device) - - # Create the stencil matrices with Vandermonde scaling - psi_s = build_stencil_s(v_s, self.order, dX.dtype) - psi_t = build_stencil_t(v_t, self.order, dY.dtype) - exponents = self.exponents.to(dX.device) - return compute_gram_entry(dX, dY, v_s, v_t, psi_s, psi_t, diagonal_count, diagonal_batch_size, longest_diagonal, ic, indices, exponents, order=self.order) - # @torch.compile(mode="max-autotune", fullgraph=True) - # def chunked_compute_signature_kernel(self, X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor: - # """ - # Compute the signature kernel between two sets of time series. - # """ - # # Generate the stencil and Vandermonde vectors - # ds = torch.tensor([1 / dX_i.shape[0]], dtype=dX_i.dtype, device=dX_i.device) - # dt = torch.tensor([1 / dY_j.shape[0]], dtype=dY_j.dtype, device=dY_j.device) - # torch.compiler.cudagraph_mark_step_begin() - # v_s, v_t = compute_vandermonde_vectors(ds, dt, order) - # psi_s = build_stencil_s(v_s, order, dX_i.device, dX_i.dtype) - # psi_t = build_stencil_t(v_t, order, dY_j.device, dY_j.dtype) - - # diagonal_count = dX_i.shape[0] + dY_j.shape[0] - 1 - # return compute_gram_entry(dX, dY, v_s, v_t, psi_s, psi_t, diagonal_count, longest_diagonal, ic, indices, self.exponents, order=self.order) - - # TODO: Think about jitting this - def compute_gram_matrix(self, X: torch.Tensor, Y: torch.Tensor, symmetric: bool = False) -> torch.Tensor: - """ - Compute the Gram matrix between two sets of time series. - Args: - X: JAX array of shape (batch_size,length, dim) representing the first set of time series - Y: JAX array of shape (batch_size, length, dim) representing the second set of time series - symmetric: If True, computes the kernel matrix for the combined set of X and Y. Default is False. - - Returns: - A JAX array of shape (batch_size, batch_size) containing the Gram matrix between X and Y - """ - gram_matrix = jnp.zeros([X.shape[0], Y.shape[0]], dtype=X.dtype, device=X.device) - - # These will stay the same for the entire batch - ds = 1.0 / X.shape[1] - dt = 1.0 / Y.shape[1] - v_s, v_t = compute_vandermonde_vectors(ds, dt, self.order, dtype=jnp.float64) - psi_s = build_stencil_s(v_s, order=self.order, dtype=X.dtype) - psi_t = build_stencil_t(v_t, order=self.order, dtype=X.dtype) - ic = torch.zeros([self.order], dtype=X.dtype).at[0].set(1) - longest_diagonal = min(X.shape[1], Y.shape[1]) - diagonal_count = X.shape[1] + Y.shape[1] - 1 - indices = torch.arange(longest_diagonal) - - - dX = torch_compute_derivative_batch(X) - dY = torch_compute_derivative_batch(Y) - for i in range(X.shape[0]): - for j in range(Y.shape[0]): - gram_matrix[i,j] = compute_gram_entry(dX, dY, v_s, v_t, psi_s, psi_t, diagonal_count, longest_diagonal, ic, indices, self.exponents) +DIAGONAL_CHUNK_SIZE = 1024 +JIT_BOUNDARY_THRESHOLD = 64 +_MAX_BLOCK_SIZE_SMALL = 256 +_MAX_BLOCK_SIZE_LARGE = 16384 - return gram_matrix -@torch.compile(mode="max-autotune", fullgraph=True) -def batch_ADM_for_diagonal( - rho: torch.Tensor, - U_buf: torch.Tensor, - S: torch.Tensor, - T: torch.Tensor, - stencil: torch.Tensor -) -> torch.Tensor: - """ - Use ADM to compute the truncated power series representation for each tile on the diagonal with refinement determined by the shape of stencil. - Args: - rho: Tensor of shape (batch_size,) containing the rho values - U_buf: Pre-allocated buffer for U matrices of shape (max_batch_size, n, n) - S: Tensor of shape (batch_size, n) containing coefficients for diagonals 0...n-1 - T: Tensor of shape (batch_size, n) containing coefficients for diagonals 0...-(n-1) - stencil: Tensor of shape (n, n) containing the initial condition - """ - # length of current diagonal is batch_size and determined by rho - batch_size = rho.shape[0] - n = stencil.shape[0] - U = U_buf[:batch_size, :, :] - U[:] = stencil - rho = rho.view(batch_size,1) - rho_powers = rho.view(batch_size,1) ** torch.arange(n, device=rho.device, dtype=rho.dtype) - # for exponent in range(n): - # U[:, exponent, exponent+1:] *= S[:, 1:S.shape[1]-exponent] * (rho ** exponent) - # U[:, exponent:, exponent] *= T[:, :T.shape[1]-exponent] * (rho ** exponent) - - # Iterate over all diagonals from -(n-1) (bottom-left diagonal) to (n-1) (top-right diagonal) - for k in range(-(n - 1), n): - # multiply_diagonal(U, k, S, T, vandermonde_full) - - # Calculate the length of the diagonal - diag_length = n - abs(k) - - # Get the view of the diagonal for all matrices in the batch - diagonal_view = torch.diagonal(U, offset=k, dim1=1, dim2=2) - - # Take the appropriate slice of the full Vandermonde matrix - rho_diag = rho_powers[:, :diag_length] - - # Get the coefficient and reshape for broadcasting - if k > 0: - # Use S for upper diagonals (k > 0) - # Map k to index in S (1 to n-1) - # coefficients = S[:, k].view(batch_size, 1) - diagonal_view.mul_(S[:,k].view(batch_size, 1)) - else: - # Use T for main and lower diagonals (k <= 0) - # Map k to index in T (0 to n-1) - # coefficients = T[:, -k].view(batch_size, 1) - diagonal_view.mul_(T[:, -k].view(batch_size, 1)) +def get_diagonal_range_bool(d: int, rows: int, cols: int) -> Tuple[int, int, int]: + # Benchmark-only version matching the corrected boolean arithmetic used in JAX. + t_start = (d < rows) * 0 + (d >= rows) * (d - rows + 1) + s_start = (d < rows) * d + (d >= rows) * (rows - 1) + dlen = min(s_start + 1, cols - t_start) + return int(s_start), int(t_start), int(dlen) + + +def get_max_block_size(device: torch.device) -> int: + if device.type != "cuda": + return _MAX_BLOCK_SIZE_SMALL + + try: + props = torch.cuda.get_device_properties(_device_index(device)) + if props.total_memory > 24 * 1024**3: + return _MAX_BLOCK_SIZE_LARGE + except Exception: + pass + + return _MAX_BLOCK_SIZE_SMALL + + +def estimate_bytes_per_pair(longest_diagonal: int, order: int, dtype: torch.dtype) -> int: + elem_bytes = torch.empty((), dtype=dtype).element_size() + buffers = 2 * longest_diagonal * order + toeplitz = 3 * longest_diagonal * order * order + aux = 5 * longest_diagonal * order + return elem_bytes * (buffers + toeplitz + aux) + + +def get_available_gpu_memory(device: torch.device) -> int: + if device.type != "cuda": + return 4 * 1024**3 + + try: + free, _ = torch.cuda.mem_get_info(_device_index(device)) + return int(free) + except Exception: + pass - # In-place multiplication: diagonal * coefficient * vandermonde_slice - diagonal_view.mul_(rho_diag) + try: + props = torch.cuda.get_device_properties(_device_index(device)) + return int(props.total_memory) + except Exception: + return 8 * 1024**3 - return U + +def _round_to_power_of_2(n: int) -> int: + if n <= 1: + return 1 + return 1 << (n - 1).bit_length() + + +def compute_block_size( + longest_diagonal: int, + order: int, + dtype: torch.dtype, + device: torch.device, + total_pairs: int, + safety_factor: float = 0.7, +) -> int: + max_bs = get_max_block_size(device) + per_pair = estimate_bytes_per_pair(longest_diagonal, order, dtype) + available = get_available_gpu_memory(device) + budget = int(available * safety_factor) + + if per_pair > 0: + raw = max(1, budget // per_pair) + else: + raw = max_bs + + raw = min(raw, total_pairs, max_bs) + return _round_to_power_of_2(raw) -@torch.compile(mode="max-autotune", fullgraph=True) def compute_vandermonde_vectors( - ds: float, dt: float, n: int, dtype: torch.dtype, device: torch.device + ds: float, + dt: float, + n: int, + dtype: torch.dtype = torch.float64, + device: Optional[torch.device] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - powers = torch.arange(n, device=device, dtype=dtype) - v_s = ds**powers - v_t = dt**powers + powers = torch.arange(n, dtype=dtype, device=device) + v_s = torch.pow(torch.as_tensor(ds, dtype=dtype, device=device), powers) + v_t = torch.pow(torch.as_tensor(dt, dtype=dtype, device=device), powers) return v_s, v_t -# @torch.compile(mode="max-autotune", fullgraph=True) -@torch.compile() -def build_stencil( - order: int = 32, device: torch.device=torch.device("cpu"), dtype: torch.dtype = torch.float64 +def build_psi_stencil( + order: int, + delta: float = 1.0, + dtype: torch.dtype = torch.float64, + device: Optional[torch.device] = None, ) -> torch.Tensor: - stencil = torch.ones([order, order], dtype=dtype, device=device) + psi = torch.zeros((order, order), dtype=dtype, device=device) + psi[0, :] = 1.0 - # Fill in the rest of the matrix with 1/(i*j) - i_indices = torch.arange(1, order, device=device).reshape(-1, 1) - j_indices = torch.arange(1, order, device=device).reshape(1, -1) + i_indices = torch.arange(1, order, dtype=dtype, device=device).reshape(-1, 1) + j_indices = torch.arange(order, dtype=dtype, device=device).reshape(1, -1) + psi[1:, :] = 1.0 / (i_indices * (j_indices + i_indices)) - stencil[1:, 1:] /= i_indices - stencil[1:, 1:] /= j_indices + powers = torch.arange(order, dtype=dtype, device=device) + psi[0, :] *= delta**powers + for row_idx in range(1, order): + psi[row_idx, :] = psi[row_idx, :] * psi[row_idx - 1, :] * delta - # Replace each diagonal with its cumulative product - for k in range(-(order - 1), order): - diag = torch.diagonal(stencil, offset=k) - diag[:] = torch.cumprod(diag, dim=0) + return psi - return stencil -# @torch.compile(mode="max-autotune", fullgraph=True,disable=False) -@torch.compile() -def build_stencil_t(v_t: torch.Tensor, order: int = 32, device: torch.device = None, dtype: torch.dtype = torch.float64) -> torch.Tensor: - """ - Build stencil matrix and multiply each row by v_t in place. - - Args: - v_t: Vandermonde vector for t direction of shape (order,) - order: Order of the polynomial approximation - device: Device to create tensor on - dtype: Data type of the tensor - - Returns: - Stencil matrix with columns multiplied by v_t - """ - # First build the standard stencil - stencil = build_stencil(order=order, device=device, dtype=dtype) - - # Multiply each column by v_t in place - # Since v_t has the same length as the columns, we can use broadcasting - # by indexing with None/newaxis along the row dimension - stencil.mul_(v_t.view(-1, 1)) - - return stencil +def build_stencil( + order: int = 32, + dtype: torch.dtype = torch.float64, + device: Optional[torch.device] = None, +) -> torch.Tensor: + stencil = torch.ones((order, order), dtype=dtype, device=device) + + i_indices = torch.arange(1, order, dtype=dtype, device=device).reshape(-1, 1) + j_indices = torch.arange(1, order, dtype=dtype, device=device).reshape(1, -1) + stencil[1:, 1:] = 1.0 / (i_indices * j_indices) + + for k in range(-(order - 1), order): + diagonal = torch.diagonal(stencil, offset=k) + diagonal.copy_(torch.cumprod(diagonal, dim=0)) -# @torch.compile(mode="max-autotune", fullgraph=True,disable=False) -@torch.compile() -def build_stencil_s(v_s: torch.Tensor, order: int = 32, device: torch.device = None, dtype: torch.dtype = torch.float64) -> torch.Tensor: - """ - Build stencil matrix and multiply each column by v_s in place. - - Args: - v_s: Vandermonde vector for s direction of shape (order,) - order: Order of the polynomial approximation - device: Device to create tensor on - dtype: Data type of the tensor - - Returns: - Stencil matrix with rows multiplied by v_s - """ - # First build the standard stencil - stencil = build_stencil(order=order, device=device, dtype=dtype) - - # Multiply each row by v_s in place - stencil.mul_(v_s) - return stencil -@torch.compile(mode="max-autotune", fullgraph=True) -def batch_compute_boundaries( - U: torch.Tensor, - S_buf: torch.Tensor, - T_buf: torch.Tensor, - v_s: torch.Tensor, - v_t: torch.Tensor, - skip_first: bool = False, - skip_last: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Compute the boundary tensor power series for a given diagonal. - - Args: - U: Tensor of shape (batch_size, n, n) containing the power series coefficients - S_buf: Pre-allocated buffer for S of shape (max_batch_size, n) - T_buf: Pre-allocated buffer for T of shape (max_batch_size, n) - v_s: Vandermonde vector for s direction - v_t: Vandermonde vector for t direction - skip_first: Whether to skip propagating the rightmost boundary of the first tile in the diagonal - skip_last: Whether to skip propagating the topmost boundary of the last tile of the diagonal - """ - - # Diagonal will always grow until it reaches top (skip_last) or right (skip_first) of grid - if skip_first and skip_last: - # Shrinking - next_dlen = U.shape[0] - 1 - # T = torch.empty((next_dlen, U.shape[1]), dtype=U.dtype, device=U.device) - # S = torch.empty((next_dlen, U.shape[1]), dtype=U.dtype, device=U.device) - T = T_buf[:next_dlen, :] - S = S_buf[:next_dlen, :] - - # Skip first, don't propagate coefficients right - torch.matmul(U[1:, :, :], v_s, out=T) - # Skip last, don't propagate coefficients up - torch.matmul(v_t, U[:-1, :, :], out=S) - - elif not skip_first and not skip_last: - # Growing - next_dlen = U.shape[0] + 1 - T = T_buf[:next_dlen, :] - S = S_buf[:next_dlen, :] - - # Top tile already has initial left boundary, tiles below propagate top boundary - torch.matmul(U, v_s, out=T[:-1, :]) - - # Bottom tile already has initial bottom boundary, tiles above propagate right boundary - torch.matmul(v_t, U, out=S[1:, :]) - elif skip_first and not skip_last: - # Staying the same size - next_dlen = U.shape[0] - T = T_buf[:next_dlen, :] - S = S_buf[:next_dlen, :] - - # Bottom tile not propagating right boundary, but top tile receives initial left boundary - torch.matmul(v_t, U, out=S) - torch.matmul(U[1:, :, :], v_s, out=T[:-1, :]) - else: - # Staying the same size - next_dlen = U.shape[0] - T = T_buf[:next_dlen, :] - S = S_buf[:next_dlen, :] - # Top tile not propagating top boundary, but bottom tile receives initial bottom boundary - torch.matmul(v_t, U[:-1, :, :], out=S[1:, :]) - torch.matmul(U, v_s, out=T) - - return S, T - - -# @torch.compile(mode="max-autotune", fullgraph=True,disable=False) -@torch.compile(dynamic=True) -def compute_boundary( - psi_s: torch.Tensor, - psi_t: torch.Tensor, - S: torch.Tensor, - T: torch.Tensor, - rho: torch.Tensor, -): - """ - Compute the boundary tensor power series for a fixed-size chunk. - - Args: - U_s: Fixed-size chunk from larger preallocated U buffer - U_t: Fixed-size chunk from larger preallocated U buffer - S: Tensor of shape (batch_size, n) containing coefficients for upper diagonals - T: Tensor of shape (batch_size, n) containing coefficients for main and lower diagonals - rho: Tensor of shape (batch_size,) containing the rho values - offset: Offset in the larger buffer - """ - # assert psi_s.shape[0] == psi_t.shape[0], f"psi_s and psi_t must have the same batch size, but got {psi_s.shape[0]} and {psi_t.shape[0]}" - # assert S.shape[1] == psi_s.shape[1], f"S must have the same number of elements as psi_s and psi_t have columns {S.shape[0]} and {psi_s.shape[1]}" - # assert T.shape[1] == psi_s.shape[0], f"T must have the same number of elements as psi_s and psi_t have rows {T.shape[0]} and {psi_s.shape[0]}" - - n = psi_s.shape[0] - batch_size = rho.shape[0] - U_s = psi_s.repeat(batch_size, 1, 1) - U_t = psi_t.repeat(batch_size, 1, 1) - - # rho_powers = rho.view(batch_size,1) ** torch.arange(n, device=rho.device, dtype=rho.dtype) - # Initialize U_s and U_t from batch_size tilings of psi_s and psi_t - # Use repeat for actual memory allocation since we'll modify these tensors in-place - - rho = rho.view(batch_size,1) - - for exponent in range(n): - rho_power = rho ** exponent - s = S[:, 1:S.shape[1]-exponent] - t = T[:, :T.shape[1]-exponent] - U_s[:, exponent, exponent+1:] *= s - U_s[:, exponent, exponent+1:] *= rho_power - U_s[:, exponent:, exponent] *= t - U_s[:, exponent:, exponent] *= rho_power - - U_t[:, exponent, exponent+1:] *= s - U_t[:, exponent, exponent+1:] *= rho_power - U_t[:, exponent:, exponent] *= t - U_t[:, exponent:, exponent] *= rho_power - - # Iterate over all diagonals from -(n-1) (bottom-left diagonal) to (n-1) (top-right diagonal) - # for k in range(-(n - 1), 1): - # diag_index = -k - # diag_length = n - diag_index - # diagonals_of_U_s = torch.diagonal(U_s, offset=k, dim1=1, dim2=2) - # diagonals_of_U_s.mul_(T[:, diag_index].view(batch_size,1)) - # diagonals_of_U_s.mul_(rho_powers[:,:diag_length]) - - # diagonals_of_U_t = torch.diagonal(U_t, offset=k, dim1=1, dim2=2) - # diagonals_of_U_t.mul_(T[:, diag_index].view(batch_size,1)) - # diagonals_of_U_t.mul_(rho_powers[:,:diag_length]) - - # for k in range(1, n): - # diag_index = k - # diag_length = n - diag_index - # diagonals_of_U_s = torch.diagonal(U_s, offset=k, dim1=1, dim2=2) - # diagonals_of_U_s.mul_(S[:, diag_index].view(batch_size,1)) - # diagonals_of_U_s.mul_(rho_powers[:,:diag_length])n - - # diagonals_of_U_t = torch.diagonal(U_t, offset=k, dim1=1, dim2=2) - # diagonals_of_U_t.mul_(S[:, diag_index].view(batch_size,1)) - # diagonals_of_U_t.mul_(rho_powers[:,:diag_length]) - - # sum cols, sum rows - return U_t.sum(dim=1), U_s.sum(dim=2) - - -def compute_boundary_inplace(psi_s: torch.Tensor, psi_t: torch.Tensor, exponents: torch.Tensor, S: torch.Tensor, T: torch.Tensor, rho: torch.Tensor): - """ - Compute the boundary tensor power series for a fixed-size chunk. - - Args: - psi_s: Fixed-size chunk from larger preallocated U buffer - psi_t: Fixed-size chunk from larger preallocated U buffer - S: Tensor of shape (n) containing coefficients for upper diagonals - T: Tensor of shape (n) containing coefficients for main and lower diagonals - rho: Tensor of shape (batch_size,) containing the rho values - offset: Offset in the larger buffer - """ - U = rho ** exponents - - def toeplitz(index): - torch.diagonal(U, offset=index, dim1=1, dim2=2)[:,index:].mul_( (index == 0) * S[index] + (index != 0) * T[index]) - torch.diagonal(U, offset=index, dim1=1, dim2=2)[:,index:].mul_(S[index]) - U[:,index:] = T[:T.shape[0]-index] - - torch.vmap(toeplitz,out_dims=None)(exponents[-1]) - - # Use direct broadcasting for element-wise multiplication - # JAX will automatically broadcast psi_s and psi_t [n, n] to match U [batch_size, n, n] - U_s = U * psi_s # Broadcasting happens automatically - U_t = U * psi_t # Broadcasting happens automatically - - # Sum all rows of U_s and all columns of U_t within each batch and store directly in S and T - S = torch.sum(U_t, axis=0, out=S) - T = torch.sum(U_s, axis=1, out=T) - - return S, T - -def stable_compute_boundaries(rho:torch.Tensor, psi_s: torch.Tensor, psi_t: torch.Tensor, exponents: torch.Tensor, s: torch.Tensor, t: torch.Tensor): - """ - Compute the boundary tensor power series for a fixed-size chunk. - - Args: - psi_s: Fixed-size chunk from larger preallocated U buffer - psi_t: Fixed-size chunk from larger preallocated U buffer - S: Tensor of shape (n) containing coefficients for upper diagonals - T: Tensor of shape (n) containing coefficients for main and lower diagonals - rho: Tensor of shape (batch_size,) containing the rho values - offset: Offset in the larger buffer - """ - # U = rho ** exponents - - def outer(i,tv,exprow): - def inner(j,sv,exp): - r = rho ** (((i <= j) * i ) + ((i>j) * j)) - return r * ((tv * (j<=i)) + (sv * (j>i) )) - return torch.vmap(inner,)(exponents[-1], s, exprow) - - U = torch.vmap(outer)(exponents[-1], t, exponents) - - # Use direct broadcasting for element-wise multiplication - # JAX will automatically broadcast psi_s and psi_t [n, n] to match U [batch_size, n, n] - U_s = U * psi_s # Broadcasting happens automatically - U_t = U * psi_t # Broadcasting happens automatically - - # Sum all rows of U_s and all columns of U_t within each batch and store directly in S and T - S = torch.sum(U_t, axis=0) - T = torch.sum(U_s, axis=1) - - return S, T - -def map_diagonal_entry(dX_i, dY_j, psi_s, psi_t,exponents, s_coeff, t_coeff, s_start: int, t_start: int, diagonal_index: int): - # Compute dot products for valid entries - rho = torch.dot(dX_i[s_start - diagonal_index,], dY_j[t_start+ diagonal_index,]) - - # Process valid entries with compute_boundary - s, t = compute_boundary(psi_s, psi_t, exponents, s_coeff, t_coeff, rho) - - -@torch.compile(mode="max-autotune", fullgraph=True,dynamic=True) -def compute_gram_entry_vmap( - dX_i: torch.Tensor, - dY_j: torch.Tensor, + +def build_stencil_s( v_s: torch.Tensor, - v_t: torch.Tensor, - psi_s: torch.Tensor, - psi_t: torch.Tensor, - diagonal_count: int, - longest_diagonal: int, - ic: torch.Tensor, - indices: torch.Tensor, - exponents: torch.Tensor, order: int = 32, + dtype: torch.dtype = torch.float64, + device: Optional[torch.device] = None, ) -> torch.Tensor: - """ - Compute the gram matrix entry using a batched approach. - - Args: - dX_i: First time series derivatives - dY_j: Second time series derivatives - v_s: Vandermonde vector for s direction - v_t: Vandermonde vector for t direction - psi_s: First time series power series coefficients - psi_t: Second time series power series coefficients - diagonal_count: Number of diagonals to compute - longest_diagonal: Longest diagonal to compute - ic: Initial condition for the power series - indices: Indices of the diagonals to compute - exponents: Exponents of the power series - order: Order of the polynomial approximation - - Returns: - Gram matrix entry (scalar) - """ - # Initialize buffers with proper shapes - S_buf = torch.zeros([longest_diagonal, order], dtype=dX_i.dtype, device=dX_i.device) - T_buf = torch.zeros([longest_diagonal, order], dtype=dX_i.dtype, device=dX_i.device) - - # Initialize first elements with 1.0 - S_buf[:, 0] = 1.0 - T_buf[:, 0] = 1.0 - - for d in range(diagonal_count): - rows = dX_i.shape[0] - cols = dY_j.shape[0] - t_start = (d=cols)*(d-cols +1) - s_start = (d=cols)*(cols - 1) - dlen = min(rows - t_start, s_start + 1) - - def next_diagonal_entry(diagonal_index): - # Combine the first two where statements into a single mask - is_before_wrap = d < dX_i.shape[0] - s_index = diagonal_index - is_before_wrap - t_index = diagonal_index + (1 - is_before_wrap) - - # Avoid branching - - - # s = ((t_start + diagonal_index == 0).cuda() * ic) + ((t_start + diagonal_index != 0).cuda() * S_buf[s_index]) - # t = ((s_start - diagonal_index == 0).cuda() * ic) + ((s_start - diagonal_index != 0).cuda() * T_buf[t_index]) - # Use vectorized operations instead of control flow - s_mask = (t_start + diagonal_index == 0).to(dtype=dX_i.dtype, device=dX_i.device) - t_mask = (s_start - diagonal_index == 0).to(dtype=dX_i.dtype, device=dX_i.device) - - # jax.debug.print(""" - # d = {}, - # diagonal_index {}: - # s_start = {} - # t_start = {} - # dlen = {} - # is_before_wrap = {} - # s_index = {} - # t_index = {} - # s = {} - # t = {} - # """, d, diagonal_index, s_start, t_start, dlen, is_before_wrap, s_index, t_index, s, t) - s = s_mask * ic + (1 - s_mask) * S_buf[s_index] - t = t_mask * ic + (1 - t_mask) * T_buf[t_index] - map_diagonal_entry(dX_i, dY_j, psi_s, psi_t, exponents, s, t, s_start, t_start, diagonal_index) - - torch.vmap(next_diagonal_entry, out_dims=None)(indices[:d+1]) - - return S_buf[0] @ v_s - - - - -# @torch.compile(mode="max-autotune-no-cudagraphs",dynamic=True) -# @torch.compile(dynamic=True) -@torch.compile(mode="max-autotune", fullgraph=True,dynamic=True) -def compute_gram_entry( - dX_i: torch.Tensor, - dY_j: torch.Tensor, - v_s: torch.Tensor, + return build_stencil(order=order, dtype=dtype, device=device) * v_s + + +def build_stencil_t( v_t: torch.Tensor, - psi_s: torch.Tensor, - psi_t: torch.Tensor, - diagonal_count: int, - diagonal_batch_size: int, - longest_diagonal: int, - ic: torch.Tensor, - indices: torch.Tensor, - exponents: torch.Tensor, order: int = 32, + dtype: torch.dtype = torch.float64, + device: Optional[torch.device] = None, ) -> torch.Tensor: - # Initial tile - S_buf = torch.zeros( - [longest_diagonal+1, order], - dtype=dX_i.dtype, - device=dX_i.device, - ) - - T_buf = torch.zeros( - [longest_diagonal, order], - dtype=dX_i.dtype, - device=dX_i.device, - ) - - S_buf[:, 0] = 1 - T_buf[:, 0] = 1 - - cols = dY_j.shape[0] - rows = dX_i.shape[0] - - for d in range(0,diagonal_count, diagonal_batch_size): - s_start, t_start, dlen = get_diagonal_range(d, dX_i.shape[0], dY_j.shape[0]) - skip_first = (s_start + 1) >= dX_i.shape[0] - skip_last = (t_start + dlen) >= dY_j.shape[0] - max_diag = min(diagonal_count, d + diagonal_batch_size) - - # This length of the longest diagonal length we will get to for this unrolled piece of the loop. - if (d+1) >= longest_diagonal and max_diag <= max(rows,cols): - batch_longest_diag = longest_diagonal - elif max_diag < longest_diagonal: - batch_longest_diag = max_diag + return build_stencil(order=order, dtype=dtype, device=device) * v_t[:, None] + + +def _device_index(device: torch.device) -> int: + if device.index is not None: + return device.index + return torch.cuda.current_device() + + +class PowerSigTorch: + def __init__( + self, + order: int = 32, + static_kernel: Callable = static_kernels.linear_kernel, + device: Optional[torch.device] = None, + dtype: torch.dtype = torch.float64, + compile_forward: bool = False, + ): + self.order = order + self.dtype = dtype + self.static_kernel = static_kernel + self.compile_forward = compile_forward + + if device is None: + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") else: - batch_longest_diag = longest_diagonal - ((d + 1) - max(rows,cols)) - - diagonal_indices = indices[:batch_longest_diag] - - - print(f"d = {d}") - print(f"max_diag = {max_diag}") - print(f"diagonal_count = {diagonal_count}") - # Process each chunk of the diagonal - for diagonal_index in range(d,max_diag): - t_start = (diagonal_index=cols)*(diagonal_index-cols +1) - s_start = (diagonal_index=cols)*(cols - 1) - is_before_wrap = diagonal_index < dX_i.shape[0] - - # 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)) - def next_diagonal_entry(index_in_diagonal, dx,dy,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 - # rho = torch.dot(dX_i[s_start-index_in_diagonal], dY_j[t_start+index_in_diagonal, :]) - s = ((t_start + index_in_diagonal == 0) * ic) + ((t_start + index_in_diagonal != 0) * S) - t = ((s_start - index_in_diagonal == 0) * ic) + ((s_start - index_in_diagonal != 0) * T) - - # Use torch.dot directly instead of creating intermediate arrays - # dX_idx = (s_start - index_in_diagonal) * ((s_start - index_in_diagonal) < dX_i.shape[0]) - # dY_idx = (t_start + index_in_diagonal) * ((t_start + index_in_diagonal) < dY_j.shape[0]) - rho = torch.dot(dx, dy) - - return stable_compute_boundaries(rho, psi_s, psi_t, exponents, s, t) - - S_next, T_next = torch.vmap(next_diagonal_entry, in_dims=(0,0,0,0,0))( - diagonal_indices, - dX_i[(s_start - diagonal_indices) * ((s_start - diagonal_indices) < dX_i.shape[0])], - dY_j[(t_start + diagonal_indices) * ((t_start + diagonal_indices) < dY_j.shape[0])], - S_buf[:diagonal_indices.shape[0]], - T_buf[:batch_longest_diag]) - - print(f"S_next = {S_next}") - print(f"T_next = {T_next}") - print(f"is_before_wrap = {is_before_wrap}") - print(f"diagonal_index = {diagonal_index}") - S_buf[is_before_wrap:S_next.shape[0]+is_before_wrap], T_buf[:T_next.shape[0]] = S_next, T_next[(1-is_before_wrap):] - - print(f"S_buf = {S_buf}") - print(f"T_buf = {T_buf}") - return S_buf[0] @ v_s - -@torch.compile(mode="max-autotune", fullgraph=True,dynamic=True) -def batch_compute_gram_entry_psi( - dX_i: torch.Tensor, - dY_j: torch.Tensor, - order: int = 32, -) -> torch.Tensor: - # Preprocessing - dX_i[:] = dX_i.flip(0) - longest_diagonal = min(dX_i.shape[0], dY_j.shape[0]) - - # Initial tile - S_buf = torch.zeros([longest_diagonal+1, order], dtype=dX_i.dtype, device=dX_i.device) - T_buf = torch.zeros([longest_diagonal+1, order], dtype=dX_i.dtype, device=dX_i.device) - S_buf[:, 0] = 1 - T_buf[:, 0] = 1 - - - - # Generate the stencil and Vandermonde vectors - ds = torch.tensor([1 / dX_i.shape[0]], dtype=dX_i.dtype, device=dX_i.device) - dt = torch.tensor([1 / dY_j.shape[0]], dtype=dY_j.dtype, device=dY_j.device) - torch.compiler.cudagraph_mark_step_begin() - v_s, v_t = compute_vandermonde_vectors(ds, dt, order) - psi_s = build_stencil_s(v_s, order, dX_i.device, dX_i.dtype) - psi_t = build_stencil_t(v_t, order, dY_j.device, dY_j.dtype) - - diagonal_count = dX_i.shape[0] + dY_j.shape[0] - 1 - - for d in range(diagonal_count): - s_start, t_start, dlen = get_diagonal_range(d, dX_i.shape[0], dY_j.shape[0]) - S = S_buf[t_start:t_start+dlen, :] - T = T_buf[t_start:t_start+dlen, :] - - dX_L = dX_i.shape[0] - (s_start + 1) - - rho = torch_compute_dot_prod_batch( - dX_i[dX_L : dX_L + dlen], - dY_j[t_start : (t_start + dlen)], + self.device = torch.device(device) + + self.exponents = torch.arange(self.order, dtype=torch.int64, device=self.device) + self.psi_s = build_psi_stencil(self.order, dtype=self.dtype, device=self.device) + for i in range(1, self.order): + self.psi_s[i, -i:] = 0.0 + + self.psi_t = build_stencil(self.order, dtype=self.dtype, device=self.device) + self.psi_s_t = self.psi_s.transpose(0, 1).contiguous() + self.psi_t_upper = torch.triu(self.psi_t, diagonal=1).contiguous() + self.psi_t_lower = torch.tril(self.psi_t, diagonal=-1).contiguous() + + self.ic = torch.zeros(self.order, dtype=self.dtype, device=self.device) + self.ic[0] = 1.0 + self.v_s_unit, self.v_t_unit = compute_vandermonde_vectors( + 1.0, 1.0, self.order, dtype=self.dtype, device=self.device ) - S_next, T_next = compute_boundary(psi_s, psi_t, S, T, rho) - - S_buf[t_start+1:t_start+dlen+1,:] = S_next - T_buf[t_start:t_start+dlen,:] = T_next - - if dlen == 1 and d == diagonal_count - 1: - return v_t @ S_buf[t_start+1] + indices = torch.arange(self.order, dtype=torch.long, device=self.device) + diff = indices[None, :] - indices[:, None] + self._toeplitz_is_upper = (diff >= 0).contiguous() + self._toeplitz_upper_index = diff.clamp_min(0).contiguous() + self._toeplitz_lower_index = (-diff).clamp_min(0).contiguous() + jac_row_idx = torch.arange(self.order, dtype=torch.long, device=self.device)[:, None] + jac_col_idx = torch.arange(self.order, dtype=torch.long, device=self.device)[None, :] + jac_diff = jac_row_idx - jac_col_idx + self._jac_row_idx = jac_row_idx.expand(self.order, self.order).contiguous() + self._jac_safe_diff = jac_diff.clamp_min(0).contiguous() + self._jac_mask = ((jac_col_idx >= 1) & (jac_row_idx >= jac_col_idx)).contiguous() + + self._minimum_sweep = self._compute_gram_entry_batch_minimum + self._boolean_sweep = self._compute_gram_entry_batch_bool + if self.device.type == "cuda" and self.compile_forward: + try: + self._minimum_sweep = torch.compile( + self._compute_gram_entry_batch_minimum, + dynamic=False, + mode="reduce-overhead", + ) + self._boolean_sweep = torch.compile( + self._compute_gram_entry_batch_bool, + dynamic=False, + mode="reduce-overhead", + ) + except Exception: + self._minimum_sweep = self._compute_gram_entry_batch_minimum + self._boolean_sweep = self._compute_gram_entry_batch_bool + + def __call__( + self, + X, + Y=None, + symmetric: bool = False, + block_size: Optional[int] = None, + show_progress: bool = True, + ) -> torch.Tensor: + return self.compute_gram_matrix( + X, Y, symmetric=symmetric, block_size=block_size, show_progress=show_progress + ) + def compute_signature_kernel( + self, X: torch.Tensor, Y: torch.Tensor, device: Optional[torch.device] = None + ) -> torch.Tensor: + X_i = self._as_single_path(X, device=device) + Y_j = self._as_single_path(Y, device=device) + return self._minimum_sweep(X_i[None, ...], Y_j[None, ...])[0] + + def compute_signature_kernel_bool_geometry( + self, X: torch.Tensor, Y: torch.Tensor, device: Optional[torch.device] = None + ) -> torch.Tensor: + X_i = self._as_single_path(X, device=device) + Y_j = self._as_single_path(Y, device=device) + return self._boolean_sweep(X_i[None, ...], Y_j[None, ...])[0] + + def compute_signature_kernel_chunked( + self, X: torch.Tensor, Y: torch.Tensor, device: Optional[torch.device] = None + ) -> torch.Tensor: + # Unlike the JAX backend, the Torch sweep is already shape-stable and does + # not need a separate long-diagonal implementation. + return self.compute_signature_kernel(X, Y, device=device) + + def compute_gram_matrix( + self, + X, + Y=None, + symmetric: bool = False, + block_size: Optional[int] = None, + show_progress: bool = True, + ) -> torch.Tensor: + X_batch = self._as_path_batch(X) + Y_batch = X_batch if Y is None else self._as_path_batch(Y) + + gram_matrix = torch.zeros( + (X_batch.shape[0], Y_batch.shape[0]), dtype=self.dtype, device=self.device + ) -def batch_compute_gram_entry( - dX_i: torch.Tensor, - dY_j: torch.Tensor, - order: int = 32, -) -> torch.Tensor: - # Preprocessing - dX_i[:] = dX_i.flip(0) - longest_diagonal = min(dX_i.shape[0], dY_j.shape[0]) - torch.compiler.cudagraph_mark_step_begin() - stencil = build_stencil(order, dX_i.device, dX_i.dtype) - # Initial tile - u_buf = torch.empty( - [longest_diagonal, stencil.shape[0], stencil.shape[1]], - dtype=dX_i.dtype, - device=dX_i.device, - ) - S_buf = torch.zeros([longest_diagonal, order], dtype=dX_i.dtype, device=dX_i.device) - T_buf = torch.zeros([longest_diagonal, order], dtype=dX_i.dtype, device=dX_i.device) - S_buf[:, 0] = 1 - T_buf[:, 0] = 1 - - u = u_buf[:1, :, :] - S = S_buf[:1, :] - T = T_buf[:1, :] - - - # Generate the stencil and Vandermonde vectors - ds = torch.tensor([1 / dX_i.shape[0]], dtype=dX_i.dtype, device=dX_i.device) - dt = torch.tensor([1 / dY_j.shape[0]], dtype=dY_j.dtype, device=dY_j.device) - v_s, v_t = compute_vandermonde_vectors(ds, dt, order) - - diagonal_count = dX_i.shape[0] + dY_j.shape[0] - 1 - - for d in range(diagonal_count): - s_start, t_start, dlen = get_diagonal_range(d, dX_i.shape[0], dY_j.shape[0]) - - dX_L = dX_i.shape[0] - (s_start + 1) - # print(f"dX_L = {dX_L}") - # print(f"s_start = {s_start}") - rho = torch_compute_dot_prod_batch( - dX_i[dX_L : dX_L + dlen], - dY_j[t_start : (t_start + dlen)], + rows = X_batch.shape[1] - 1 + cols = Y_batch.shape[1] - 1 + if rows <= 0 or cols <= 0: + raise ValueError("paths must have length at least 2") + + pairs_i = [] + pairs_j = [] + for i in range(X_batch.shape[0]): + for j in range(i if symmetric else 0, Y_batch.shape[0]): + pairs_i.append(i) + pairs_j.append(j) + + total_pairs = len(pairs_i) + if total_pairs == 0: + return gram_matrix + + longest_diagonal = min(rows, cols) + if block_size is None: + block_size = compute_block_size( + longest_diagonal, self.order, self.dtype, self.device, total_pairs + ) + else: + block_size = _round_to_power_of_2(min(block_size, total_pairs)) + + i_all = torch.tensor(pairs_i, dtype=torch.long, device=self.device) + j_all = torch.tensor(pairs_j, dtype=torch.long, device=self.device) + + pbar = tqdm(total=total_pairs, desc="Computing Gram Matrix", disable=not show_progress) + offset = 0 + while offset < total_pairs: + end = min(offset + block_size, total_pairs) + batch_i = i_all[offset:end] + batch_j = j_all[offset:end] + actual_count = end - offset + + if actual_count < block_size: + pad_count = block_size - actual_count + batch_i = torch.cat([batch_i, batch_i[-1:].expand(pad_count)]) + batch_j = torch.cat([batch_j, batch_j[-1:].expand(pad_count)]) + + results = self._minimum_sweep(X_batch[batch_i], Y_batch[batch_j]) + actual_i = i_all[offset:end] + actual_j = j_all[offset:end] + gram_matrix[actual_i, actual_j] = results[:actual_count] + if symmetric: + gram_matrix[actual_j, actual_i] = results[:actual_count] + + pbar.update(end - offset) + offset = end + + pbar.close() + return gram_matrix + + def compute_prefix_family( + self, + state_path, + refs, + min_prefix_len: int = 2, + max_prefix_len: Optional[int] = None, + checkpoint_interval: Optional[int] = None, + ) -> torch.Tensor: + from .autodiff import compute_prefix_family + + return compute_prefix_family( + self, + state_path, + refs, + min_prefix_len=min_prefix_len, + max_prefix_len=max_prefix_len, + checkpoint_interval=checkpoint_interval, ) - u = batch_ADM_for_diagonal(rho, u_buf, S, T, stencil) + def compute_prefix_family_fast_diff( + self, + state_path, + refs, + min_prefix_len: int = 2, + max_prefix_len: Optional[int] = None, + checkpoint_interval: Optional[int] = None, + ) -> torch.Tensor: + return self.compute_prefix_family( + state_path, + refs, + min_prefix_len=min_prefix_len, + max_prefix_len=max_prefix_len, + checkpoint_interval=checkpoint_interval, + ) + + def compute_signature_kernel_fast_diff( + self, + X, + Y, + checkpoint_interval: Optional[int] = None, + ) -> torch.Tensor: + from .autodiff import compute_sig_kernel_fast_diff + + return compute_sig_kernel_fast_diff( + self, + X, + Y, + checkpoint_interval=checkpoint_interval, + ) - if d == diagonal_count - 1: - return torch.einsum("i,ij,j->", v_t, u[0], v_s) - - skip_first = (s_start + 1) >= dX_i.shape[0] - skip_last = (t_start + dlen) >= dY_j.shape[0] + def compute_gram_fast_diff( + self, + X, + Y, + symmetric: bool = False, + block_size: Optional[int] = None, + checkpoint_interval: Optional[int] = None, + show_progress: bool = True, + ) -> torch.Tensor: + from .autodiff import compute_gram_fast_diff + + return compute_gram_fast_diff( + self, + X, + Y, + symmetric=symmetric, + block_size=block_size, + checkpoint_interval=checkpoint_interval, + show_progress=show_progress, + ) - # old_S, old_T = S, T - S, T = batch_compute_boundaries( - u, S_buf, T_buf, v_s, v_t, skip_first=skip_first, skip_last=skip_last + def _as_tensor(self, X, device: Optional[torch.device] = None) -> torch.Tensor: + target_device = self.device if device is None else torch.device(device) + if torch.is_tensor(X): + return X.to(device=target_device, dtype=self.dtype) + return torch.as_tensor(X, dtype=self.dtype, device=target_device) + + def _as_single_path(self, X, device: Optional[torch.device] = None) -> torch.Tensor: + path = self._as_tensor(X, device=device) + if path.ndim == 3: + if path.shape[0] != 1: + raise ValueError("expected a single path with shape (length, dim)") + path = path[0] + if path.ndim != 2: + raise ValueError("expected a path with shape (length, dim)") + if path.shape[0] < 2: + raise ValueError("paths must have length at least 2") + return path + + def _as_path_batch(self, X) -> torch.Tensor: + paths = self._as_tensor(X) + if paths.ndim == 2: + paths = paths[None, ...] + if paths.ndim != 3: + raise ValueError("expected paths with shape (batch, length, dim)") + if paths.shape[1] < 2: + raise ValueError("paths must have length at least 2") + return paths + + def _compute_gram_entry_batch( + self, + X_batch: torch.Tensor, + Y_batch: torch.Tensor, + geometry_fn: Callable[[int, int, int], Tuple[int, int, int]] = get_diagonal_range, + ) -> torch.Tensor: + pair_batch = X_batch.shape[0] + rows = X_batch.shape[1] - 1 + cols = Y_batch.shape[1] - 1 + longest_diagonal = min(rows, cols) + diagonal_count = rows + cols - 1 + + S_buf = torch.zeros( + (pair_batch, longest_diagonal, self.order), dtype=self.dtype, device=self.device ) - # del old_S, old_T - - # return torch.matmul(torch.matmul(v_t, u), v_s).item() - # return torch.einsum("i,bij,j->", v_t, u, v_s) - - -@torch.compile(mode="max-autotune", fullgraph=True) -def build_increasing_matrix(n: int, dtype=torch.int8, device=None) -> torch.Tensor: - """ - Build an n x n matrix where each value is the maximum of its row and column indices. - For example, for n=4: - [[0, 0, 0, 0], - [0, 1, 1, 1], - [0, 1, 2, 2], - [0, 1, 2, 3]] - - Args: - n: Size of the matrix - dtype: Data type of the matrix - - Returns: - Matrix of shape (n, n) with the specified pattern - """ - # Create row and column indices - rows = torch.arange(n, dtype=dtype, device=device)[:, None] # Shape: (n, 1) - cols = torch.arange(n, dtype=dtype, device=device)[None, :] # Shape: (1, n) - - # Take maximum of row and column indices - matrix = torch.minimum(rows, cols) - - return matrix \ No newline at end of file + T_buf = torch.zeros_like(S_buf) + S_buf[:, :, 0] = 1.0 + T_buf[:, :, 0] = 1.0 + + ic = self.ic.view(1, 1, self.order) + max_index = max(longest_diagonal - 1, 0) + + for d in range(diagonal_count): + s_start, t_start, dlen = geometry_fn(d, rows, cols) + diagonal_indices = torch.arange(dlen, dtype=torch.long, device=self.device) + row_indices = s_start - diagonal_indices + col_indices = t_start + diagonal_indices + is_before_wrap = d < rows + + s_index = torch.clamp(diagonal_indices - int(is_before_wrap), 0, max_index) + t_index = torch.clamp(diagonal_indices + int(not is_before_wrap), 0, max_index) + + s_prev = S_buf[:, s_index, :] + t_prev = T_buf[:, t_index, :] + + s = torch.where((t_start + diagonal_indices).view(1, dlen, 1) == 0, ic, s_prev) + t = torch.where((s_start - diagonal_indices).view(1, dlen, 1) == 0, ic, t_prev) + + rho = self._evaluate_static_kernel( + X_batch[:, row_indices + 1, :], + X_batch[:, row_indices, :], + Y_batch[:, col_indices + 1, :], + Y_batch[:, col_indices, :], + ) + + S_next, T_next = self._map_diagonal_entry_batch(rho, s, t) + pad_rows = longest_diagonal - dlen + S_buf = F.pad(S_next, (0, 0, 0, pad_rows)) + T_buf = F.pad(T_next, (0, 0, 0, pad_rows)) + + return torch.matmul(S_buf[:, 0, :], self.v_s_unit) + + def _compute_gram_entry_batch_minimum( + self, X_batch: torch.Tensor, Y_batch: torch.Tensor + ) -> torch.Tensor: + return self._compute_gram_entry_batch(X_batch, Y_batch, geometry_fn=get_diagonal_range) + + def _compute_gram_entry_batch_bool( + self, X_batch: torch.Tensor, Y_batch: torch.Tensor + ) -> torch.Tensor: + return self._compute_gram_entry_batch(X_batch, Y_batch, geometry_fn=get_diagonal_range_bool) + + def _evaluate_static_kernel( + self, + x2: torch.Tensor, + x1: torch.Tensor, + y2: torch.Tensor, + y1: torch.Tensor, + ) -> torch.Tensor: + try: + out = self.static_kernel(x2, x1, y2, y1) + if out.shape == x2.shape[:-1]: + return out + except Exception: + pass + + x2_flat = x2.reshape(-1, x2.shape[-1]) + x1_flat = x1.reshape(-1, x1.shape[-1]) + y2_flat = y2.reshape(-1, y2.shape[-1]) + y1_flat = y1.reshape(-1, y1.shape[-1]) + out = torch.vmap(self.static_kernel)(x2_flat, x1_flat, y2_flat, y1_flat) + return out.reshape(x2.shape[:-1]) + + def _map_diagonal_entry_batch( + self, rho: torch.Tensor, s: torch.Tensor, t: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + leading_shape = rho.shape + rho_flat = rho.reshape(-1) + s_flat = s.reshape(-1, self.order) + t_flat = t.reshape(-1, self.order) + + r = torch.pow(rho_flat[:, None], self.exponents) + toeplitz = self._build_toeplitz_batch(t_flat, s_flat) + + s_dense = torch.matmul(t_flat, self.psi_s_t) + t_dense = torch.matmul(s_flat, self.psi_s_t) + + s_next = torch.bmm(r.unsqueeze(1), toeplitz * self.psi_t_upper).squeeze(1) + s_next = s_next + (s_dense * r) + + t_next = torch.bmm(toeplitz * self.psi_t_lower, r.unsqueeze(-1)).squeeze(-1) + t_next = t_next + (t_dense * r) + + new_shape = (*leading_shape, self.order) + return s_next.reshape(new_shape), t_next.reshape(new_shape) + + def _build_toeplitz_batch(self, first_col: torch.Tensor, first_row: torch.Tensor) -> torch.Tensor: + upper = first_row[:, self._toeplitz_upper_index] + lower = first_col[:, self._toeplitz_lower_index] + return torch.where(self._toeplitz_is_upper.unsqueeze(0), upper, lower) diff --git a/powersig/torch/autodiff.py b/powersig/torch/autodiff.py new file mode 100644 index 0000000..bb6e836 --- /dev/null +++ b/powersig/torch/autodiff.py @@ -0,0 +1,503 @@ +from functools import partial +from math import ceil, sqrt +from typing import Optional, Tuple + +import torch +from tqdm.auto import tqdm + +from powersig.util.grid import get_diagonal_range + +from . import static_kernels + + +def _compute_checkpoint_interval(diagonal_count: int) -> int: + return max(1, min(int(sqrt(diagonal_count)), diagonal_count)) + + +def _extract_rbf_bandwidth(static_kernel) -> Optional[float]: + if static_kernel is static_kernels.rbf_kernel: + return 1.0 + if isinstance(static_kernel, partial) and static_kernel.func is static_kernels.rbf_kernel: + if static_kernel.keywords and "bandwidth" in static_kernel.keywords: + return float(static_kernel.keywords["bandwidth"]) + if static_kernel.args: + return float(static_kernel.args[0]) + return 1.0 + return None + + +def _prefix_step( + ps, + state_path: torch.Tensor, + refs: torch.Tensor, + S_buf: torch.Tensor, + T_buf: torch.Tensor, + d: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + rows = state_path.shape[0] - 1 + cols = refs.shape[1] - 1 + s_start, t_start, dlen = get_diagonal_range(d, rows, cols) + diag_idx = torch.arange(dlen, dtype=torch.long, device=state_path.device) + row_idx = s_start - diag_idx + col_idx = t_start + diag_idx + + ic = ps.ic.view(1, 1, ps.order) + left_prev = S_buf[:, row_idx, :] + left = torch.where((col_idx == 0).view(1, dlen, 1), ic, left_prev) + + bottom_src_idx = torch.clamp(row_idx - 1, min=0) + bottom_prev = T_buf[:, bottom_src_idx, :] + bottom = torch.where((row_idx == 0).view(1, dlen, 1), ic, bottom_prev) + + x2 = state_path[row_idx + 1].unsqueeze(0).expand(refs.shape[0], -1, -1) + x1 = state_path[row_idx].unsqueeze(0).expand_as(x2) + y2 = refs[:, col_idx + 1, :] + y1 = refs[:, col_idx, :] + + rho = ps._evaluate_static_kernel(x2, x1, y2, y1) + s_next, t_next = ps._map_diagonal_entry_batch(rho, left, bottom) + + new_S = S_buf.clone() + new_T = T_buf.clone() + new_S[:, row_idx, :] = s_next + new_T[:, row_idx, :] = t_next + return new_S, new_T + + +def _forward_prefix_family_with_checkpoints( + ps, + state_path: torch.Tensor, + refs: torch.Tensor, + min_prefix_len: int, + max_prefix_len: int, + checkpoint_interval: int, +): + rows = state_path.shape[0] - 1 + cols = refs.shape[1] - 1 + diagonal_count = rows + cols - 1 + num_prefixes = max_prefix_len - min_prefix_len + 1 + num_refs = refs.shape[0] + num_checkpoints = ceil(diagonal_count / checkpoint_interval) + + S_buf = ps.ic.view(1, 1, ps.order).expand(num_refs, rows, ps.order).clone() + T_buf = ps.ic.view(1, 1, ps.order).expand(num_refs, rows, ps.order).clone() + emissions = torch.zeros( + (num_prefixes, num_refs), dtype=state_path.dtype, device=state_path.device + ) + S_checkpoints = torch.zeros( + (num_checkpoints, num_refs, rows, ps.order), + dtype=state_path.dtype, + device=state_path.device, + ) + T_checkpoints = torch.zeros_like(S_checkpoints) + + for d in range(diagonal_count): + if d % checkpoint_interval == 0: + ckpt_idx = d // checkpoint_interval + S_checkpoints[ckpt_idx] = S_buf + T_checkpoints[ckpt_idx] = T_buf + + S_buf, T_buf = _prefix_step(ps, state_path, refs, S_buf, T_buf, d) + + emit_i = d - cols + 1 + emit_prefix_len = emit_i + 2 + emit_valid = ( + 0 <= emit_i < rows + and min_prefix_len <= emit_prefix_len <= max_prefix_len + ) + if emit_valid: + emit_idx = emit_prefix_len - min_prefix_len + emissions[emit_idx] = torch.matmul(T_buf[:, emit_i, :], ps.v_t_unit) + + return emissions, S_checkpoints, T_checkpoints + + +def _map_diagonal_entry_bwd_batch( + ps, + rho: torch.Tensor, + s: torch.Tensor, + t: torch.Tensor, + bar_s_next: torch.Tensor, + bar_t_next: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + leading_shape = rho.shape + rho_flat = rho.reshape(-1) + s_flat = s.reshape(-1, ps.order) + t_flat = t.reshape(-1, ps.order) + bar_s_flat = bar_s_next.reshape(-1, ps.order) + bar_t_flat = bar_t_next.reshape(-1, ps.order) + + r = torch.pow(rho_flat[:, None], ps.exponents) + dr = torch.zeros_like(r) + if ps.order > 1: + orders = torch.arange(1, ps.order, dtype=ps.dtype, device=ps.device) + dr[:, 1:] = orders * r[:, :-1] + + toeplitz = ps._build_toeplitz_batch(t_flat, s_flat) * ps.psi_t + upper = torch.triu(toeplitz, diagonal=1) + lower = torch.tril(toeplitz, diagonal=-1) + + psi_s_t = torch.matmul(t_flat, ps.psi_s_t) + psi_s_s = torch.matmul(s_flat, ps.psi_s_t) + + rho_term = torch.bmm(upper, bar_s_flat.unsqueeze(-1)).squeeze(-1) + rho_term = rho_term + torch.bmm(lower.transpose(1, 2), bar_t_flat.unsqueeze(-1)).squeeze(-1) + rho_term = rho_term + psi_s_t * bar_s_flat + psi_s_s * bar_t_flat + bar_rho = torch.sum(dr * rho_term, dim=-1) + + J_ss = torch.where( + ps._jac_mask.unsqueeze(0), + r[:, ps._jac_safe_diff] * ps.psi_t[ps._jac_safe_diff, ps._jac_row_idx], + torch.zeros((1, ps.order, ps.order), dtype=ps.dtype, device=ps.device), + ) + J_tt = torch.where( + ps._jac_mask.unsqueeze(0), + r[:, ps._jac_safe_diff] * ps.psi_t[ps._jac_row_idx, ps._jac_safe_diff], + torch.zeros((1, ps.order, ps.order), dtype=ps.dtype, device=ps.device), + ) + + r_bar_s_next = r * bar_s_flat + r_bar_t_next = r * bar_t_flat + bar_s = torch.bmm(J_ss.transpose(1, 2), bar_s_flat.unsqueeze(-1)).squeeze(-1) + bar_s = bar_s + torch.matmul(r_bar_t_next, ps.psi_s) + bar_t = torch.bmm(J_tt.transpose(1, 2), bar_t_flat.unsqueeze(-1)).squeeze(-1) + bar_t = bar_t + torch.matmul(r_bar_s_next, ps.psi_s) + + new_shape = (*leading_shape, ps.order) + return bar_s.reshape(new_shape), bar_t.reshape(new_shape), bar_rho.reshape(leading_shape) + + +def _static_kernel_vjp_batch( + ps, + x2: torch.Tensor, + x1: torch.Tensor, + y2: torch.Tensor, + y1: torch.Tensor, + bar_rho: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if ps.static_kernel is static_kernels.linear_kernel: + dy = y2 - y1 + dx = x2 - x1 + scale = bar_rho.unsqueeze(-1) + return scale * dy, -scale * dy, scale * dx, -scale * dx + + bandwidth = _extract_rbf_bandwidth(ps.static_kernel) + if bandwidth is not None: + inv_bw_sq = 1.0 / (bandwidth**2) + + diff_x2y2 = x2 - y2 + diff_x1y2 = x1 - y2 + diff_x2y1 = x2 - y1 + diff_x1y1 = x1 - y1 + + k_x2y2 = static_kernels.rbf_fn(diff_x2y2, bandwidth) + k_x1y2 = static_kernels.rbf_fn(diff_x1y2, bandwidth) + k_x2y1 = static_kernels.rbf_fn(diff_x2y1, bandwidth) + k_x1y1 = static_kernels.rbf_fn(diff_x1y1, bandwidth) + + grad_a_x2y2 = -diff_x2y2 * (inv_bw_sq * k_x2y2.unsqueeze(-1)) + grad_b_x2y2 = diff_x2y2 * (inv_bw_sq * k_x2y2.unsqueeze(-1)) + grad_a_x1y2 = -diff_x1y2 * (inv_bw_sq * k_x1y2.unsqueeze(-1)) + grad_b_x1y2 = diff_x1y2 * (inv_bw_sq * k_x1y2.unsqueeze(-1)) + grad_a_x2y1 = -diff_x2y1 * (inv_bw_sq * k_x2y1.unsqueeze(-1)) + grad_b_x2y1 = diff_x2y1 * (inv_bw_sq * k_x2y1.unsqueeze(-1)) + grad_a_x1y1 = -diff_x1y1 * (inv_bw_sq * k_x1y1.unsqueeze(-1)) + grad_b_x1y1 = diff_x1y1 * (inv_bw_sq * k_x1y1.unsqueeze(-1)) + + g_x2 = grad_a_x2y2 - grad_a_x2y1 + g_x1 = -grad_a_x1y2 + grad_a_x1y1 + g_y2 = grad_b_x2y2 - grad_b_x1y2 + g_y1 = -grad_b_x2y1 + grad_b_x1y1 + scale = bar_rho.unsqueeze(-1) + return scale * g_x2, scale * g_x1, scale * g_y2, scale * g_y1 + + with torch.enable_grad(): + x2_req = x2.detach().requires_grad_(True) + x1_req = x1.detach().requires_grad_(True) + y2_req = y2.detach().requires_grad_(True) + y1_req = y1.detach().requires_grad_(True) + rho = ps._evaluate_static_kernel(x2_req, x1_req, y2_req, y1_req) + grads = torch.autograd.grad( + rho, + (x2_req, x1_req, y2_req, y1_req), + grad_outputs=bar_rho, + allow_unused=False, + ) + return grads + + +def _reverse_prefix_family_with_replay( + ps, + state_path: torch.Tensor, + refs: torch.Tensor, + grad_output: torch.Tensor, + S_checkpoints: torch.Tensor, + T_checkpoints: torch.Tensor, + min_prefix_len: int, + max_prefix_len: int, + checkpoint_interval: int, +): + rows = state_path.shape[0] - 1 + cols = refs.shape[1] - 1 + diagonal_count = rows + cols - 1 + num_refs = refs.shape[0] + num_checkpoints = ceil(diagonal_count / checkpoint_interval) + + grad_state = torch.zeros_like(state_path) + grad_refs = torch.zeros_like(refs) + bar_S = torch.zeros((num_refs, rows, ps.order), dtype=ps.dtype, device=ps.device) + bar_T = torch.zeros_like(bar_S) + + for b in range(num_checkpoints - 1, -1, -1): + block_start = b * checkpoint_interval + block_end = min(block_start + checkpoint_interval, diagonal_count) + block_len = block_end - block_start + + S_buf = S_checkpoints[b] + T_buf = T_checkpoints[b] + S_tape = [] + T_tape = [] + + for d in range(block_start, block_end): + S_tape.append(S_buf) + T_tape.append(T_buf) + S_buf, T_buf = _prefix_step(ps, state_path, refs, S_buf, T_buf, d) + + for local_idx in range(block_len - 1, -1, -1): + d = block_start + local_idx + emit_i = d - cols + 1 + emit_prefix_len = emit_i + 2 + emit_valid = ( + 0 <= emit_i < rows + and min_prefix_len <= emit_prefix_len <= max_prefix_len + ) + if emit_valid: + emit_idx = emit_prefix_len - min_prefix_len + bar_T[:, emit_i, :] = bar_T[:, emit_i, :] + ( + grad_output[emit_idx].unsqueeze(-1) * ps.v_t_unit + ) + + S_saved = S_tape[local_idx] + T_saved = T_tape[local_idx] + + s_start, t_start, dlen = get_diagonal_range(d, rows, cols) + diag_idx = torch.arange(dlen, dtype=torch.long, device=ps.device) + row_idx = s_start - diag_idx + col_idx = t_start + diag_idx + + ic = ps.ic.view(1, 1, ps.order) + left_prev = S_saved[:, row_idx, :] + left = torch.where((col_idx == 0).view(1, dlen, 1), ic, left_prev) + + bottom_src_idx = torch.clamp(row_idx - 1, min=0) + bottom_prev = T_saved[:, bottom_src_idx, :] + bottom = torch.where((row_idx == 0).view(1, dlen, 1), ic, bottom_prev) + + x2 = state_path[row_idx + 1].unsqueeze(0).expand(num_refs, -1, -1) + x1 = state_path[row_idx].unsqueeze(0).expand_as(x2) + y2 = refs[:, col_idx + 1, :] + y1 = refs[:, col_idx, :] + rho = ps._evaluate_static_kernel(x2, x1, y2, y1) + + bar_s_next = bar_S[:, row_idx, :] + bar_t_next = bar_T[:, row_idx, :] + bar_left, bar_bottom, bar_rho = _map_diagonal_entry_bwd_batch( + ps, rho, left, bottom, bar_s_next, bar_t_next + ) + + bar_S_prev = bar_S.clone() + bar_T_prev = bar_T.clone() + bar_S_prev[:, row_idx, :] = 0.0 + bar_T_prev[:, row_idx, :] = 0.0 + + left_mask = col_idx > 0 + if torch.any(left_mask): + left_rows = row_idx[left_mask] + bar_S_prev[:, left_rows, :] = bar_S_prev[:, left_rows, :] + bar_left[:, left_mask, :] + + bottom_mask = row_idx > 0 + if torch.any(bottom_mask): + bottom_rows = row_idx[bottom_mask] - 1 + bar_T_prev[:, bottom_rows, :] = bar_T_prev[:, bottom_rows, :] + bar_bottom[:, bottom_mask, :] + + g_x2, g_x1, g_y2, g_y1 = _static_kernel_vjp_batch(ps, x2, x1, y2, y1, bar_rho) + grad_state[row_idx + 1] = grad_state[row_idx + 1] + g_x2.sum(dim=0) + grad_state[row_idx] = grad_state[row_idx] + g_x1.sum(dim=0) + grad_refs[:, col_idx + 1, :] = grad_refs[:, col_idx + 1, :] + g_y2 + grad_refs[:, col_idx, :] = grad_refs[:, col_idx, :] + g_y1 + + bar_S = bar_S_prev + bar_T = bar_T_prev + + return grad_state, grad_refs + + +class _PrefixFamilyFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + state_path: torch.Tensor, + refs: torch.Tensor, + ps, + min_prefix_len: int, + max_prefix_len: int, + checkpoint_interval: int, + ) -> torch.Tensor: + out, S_ckpt, T_ckpt = _forward_prefix_family_with_checkpoints( + ps, + state_path, + refs, + min_prefix_len, + max_prefix_len, + checkpoint_interval, + ) + ctx.ps = ps + ctx.min_prefix_len = min_prefix_len + ctx.max_prefix_len = max_prefix_len + ctx.checkpoint_interval = checkpoint_interval + ctx.save_for_backward(state_path, refs, S_ckpt, T_ckpt) + return out + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + state_path, refs, S_ckpt, T_ckpt = ctx.saved_tensors + grad_state, grad_refs = _reverse_prefix_family_with_replay( + ctx.ps, + state_path, + refs, + grad_output, + S_ckpt, + T_ckpt, + ctx.min_prefix_len, + ctx.max_prefix_len, + ctx.checkpoint_interval, + ) + return grad_state, grad_refs, None, None, None, None + + +def compute_prefix_family( + ps, + state_path, + refs, + min_prefix_len: int = 2, + max_prefix_len: Optional[int] = None, + checkpoint_interval: Optional[int] = None, +) -> torch.Tensor: + state_path_t = ps._as_single_path(state_path) + refs_t = ps._as_path_batch(refs) + + if min_prefix_len < 2: + raise ValueError(f"min_prefix_len must be >= 2 (got {min_prefix_len})") + + T = state_path_t.shape[0] + if max_prefix_len is None: + max_prefix_len = T + if max_prefix_len > T: + raise ValueError(f"max_prefix_len ({max_prefix_len}) > state_path length ({T})") + if max_prefix_len < min_prefix_len: + raise ValueError( + f"max_prefix_len ({max_prefix_len}) < min_prefix_len ({min_prefix_len})" + ) + if refs_t.shape[1] < 2: + raise ValueError("reference paths must have length at least 2") + + rows = state_path_t.shape[0] - 1 + cols = refs_t.shape[1] - 1 + diagonal_count = rows + cols - 1 + if checkpoint_interval is None: + checkpoint_interval = _compute_checkpoint_interval(diagonal_count) + + return _PrefixFamilyFunction.apply( + state_path_t, + refs_t, + ps, + min_prefix_len, + max_prefix_len, + checkpoint_interval, + ) + + +def compute_prefix_family_fast_diff( + ps, + state_path, + refs, + min_prefix_len: int = 2, + max_prefix_len: Optional[int] = None, + checkpoint_interval: Optional[int] = None, +) -> torch.Tensor: + return compute_prefix_family( + ps, + state_path, + refs, + min_prefix_len=min_prefix_len, + max_prefix_len=max_prefix_len, + checkpoint_interval=checkpoint_interval, + ) + + +def compute_sig_kernel_fast_diff( + ps, + X_i, + Y_j, + checkpoint_interval: Optional[int] = None, +) -> torch.Tensor: + X_i_t = ps._as_single_path(X_i) + Y_j_t = ps._as_single_path(Y_j) + out = compute_prefix_family( + ps, + X_i_t, + Y_j_t[None, ...], + min_prefix_len=X_i_t.shape[0], + max_prefix_len=X_i_t.shape[0], + checkpoint_interval=checkpoint_interval, + ) + return out[0, 0] + + +def compute_gram_fast_diff( + ps, + X, + Y, + symmetric: bool = False, + block_size: Optional[int] = None, + checkpoint_interval: Optional[int] = None, + show_progress: bool = True, +) -> torch.Tensor: + X_batch = ps._as_path_batch(X) + Y_batch = ps._as_path_batch(Y) + + pairs_i = [] + pairs_j = [] + for i in range(X_batch.shape[0]): + for j in range(i if symmetric else 0, Y_batch.shape[0]): + pairs_i.append(i) + pairs_j.append(j) + + total_pairs = len(pairs_i) + if total_pairs == 0: + return torch.zeros( + (X_batch.shape[0], Y_batch.shape[0]), dtype=ps.dtype, device=ps.device + ) + + values = [] + pbar = tqdm(total=total_pairs, desc="Computing Gram (diff)", disable=not show_progress) + for i, j in zip(pairs_i, pairs_j): + values.append( + compute_sig_kernel_fast_diff( + ps, + X_batch[i], + Y_batch[j], + checkpoint_interval=checkpoint_interval, + ) + ) + pbar.update(1) + pbar.close() + + vals = torch.stack(values) + i_all = torch.tensor(pairs_i, dtype=torch.long, device=ps.device) + j_all = torch.tensor(pairs_j, dtype=torch.long, device=ps.device) + gram = torch.zeros((X_batch.shape[0], Y_batch.shape[0]), dtype=ps.dtype, device=ps.device) + gram = gram.index_put((i_all, j_all), vals) + if symmetric: + gram = gram.index_put((j_all, i_all), vals) + return gram diff --git a/powersig/torch/static_kernels.py b/powersig/torch/static_kernels.py new file mode 100644 index 0000000..dfb7276 --- /dev/null +++ b/powersig/torch/static_kernels.py @@ -0,0 +1,26 @@ +import torch + + +def linear_kernel( + x2: torch.Tensor, x1: torch.Tensor, y2: torch.Tensor, y1: torch.Tensor +) -> torch.Tensor: + return torch.sum((x2 - x1) * (y2 - y1), dim=-1) + + +def rbf_fn(diff: torch.Tensor, bandwidth: float = 1.0) -> torch.Tensor: + sq_dist = torch.sum(diff * diff, dim=-1) + return torch.exp(-sq_dist / (2.0 * (bandwidth**2))) + + +def rbf_kernel( + x2: torch.Tensor, + x1: torch.Tensor, + y2: torch.Tensor, + y1: torch.Tensor, + bandwidth: float = 1.0, +) -> torch.Tensor: + kx2y2 = rbf_fn(x2 - y2, bandwidth) + kx1y1 = rbf_fn(x1 - y1, bandwidth) + kx2y1 = rbf_fn(x2 - y1, bandwidth) + kx1y2 = rbf_fn(x1 - y2, bandwidth) + return (kx2y2 - kx1y2) - (kx2y1 - kx1y1) diff --git a/tests/test_autodiff_torch.py b/tests/test_autodiff_torch.py new file mode 100644 index 0000000..37b8236 --- /dev/null +++ b/tests/test_autodiff_torch.py @@ -0,0 +1,126 @@ +"""Tests for the Torch checkpointed autodiff entry points.""" + +import unittest + +import numpy as np +import torch + +from powersig.torch import ( + compute_gram_fast_diff, + compute_prefix_family, + compute_prefix_family_fast_diff, + compute_sig_kernel_fast_diff, +) +from powersig.torch.algorithm import PowerSigTorch +from powersig.torch.static_kernels import linear_kernel + + +CPU = torch.device("cpu") + + +class TestFastDiffTorch(unittest.TestCase): + def setUp(self): + self.ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + + def test_single_pair_forward_matches_native(self): + rng = np.random.default_rng(0) + X = torch.tensor(rng.standard_normal((8, 2)), dtype=torch.float64, device=CPU) + Y = torch.tensor(rng.standard_normal((7, 2)), dtype=torch.float64, device=CPU) + native = self.ps.compute_signature_kernel(X, Y) + fast = compute_sig_kernel_fast_diff(self.ps, X, Y) + np.testing.assert_allclose(float(native), float(fast), rtol=1e-10, atol=1e-12) + + def test_single_pair_grad_matches_native(self): + rng = np.random.default_rng(1) + X_np = rng.standard_normal((6, 2)) + Y_np = rng.standard_normal((5, 2)) + + X_fast = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + Y_fast = torch.tensor(Y_np, dtype=torch.float64, device=CPU, requires_grad=True) + fast = compute_sig_kernel_fast_diff(self.ps, X_fast, Y_fast) + gX_fast, gY_fast = torch.autograd.grad(fast, (X_fast, Y_fast)) + + X_native = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + Y_native = torch.tensor(Y_np, dtype=torch.float64, device=CPU, requires_grad=True) + native = self.ps.compute_signature_kernel(X_native, Y_native) + gX_native, gY_native = torch.autograd.grad(native, (X_native, Y_native)) + + np.testing.assert_allclose( + gX_fast.detach().cpu().numpy(), + gX_native.detach().cpu().numpy(), + rtol=1e-5, + atol=1e-8, + ) + np.testing.assert_allclose( + gY_fast.detach().cpu().numpy(), + gY_native.detach().cpu().numpy(), + rtol=1e-5, + atol=1e-8, + ) + + def test_gram_fast_diff_matches_native(self): + rng = np.random.default_rng(2) + X = torch.tensor(rng.standard_normal((3, 7, 2)), dtype=torch.float64, device=CPU) + Y = torch.tensor(rng.standard_normal((2, 6, 2)), dtype=torch.float64, device=CPU) + native = self.ps.compute_gram_matrix(X, Y, show_progress=False) + fast = compute_gram_fast_diff(self.ps, X, Y, show_progress=False) + np.testing.assert_allclose(native.detach().cpu().numpy(), fast.detach().cpu().numpy(), rtol=1e-10, atol=1e-12) + + def test_checkpoint_intervals_give_same_grads(self): + rng = np.random.default_rng(3) + X_np = rng.standard_normal((9, 2)) + Y_np = rng.standard_normal((8, 2)) + + grads = [] + for ckpt in [1, 2, 3, 5]: + X = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + Y = torch.tensor(Y_np, dtype=torch.float64, device=CPU, requires_grad=True) + out = compute_sig_kernel_fast_diff(self.ps, X, Y, checkpoint_interval=ckpt) + gX, gY = torch.autograd.grad(out, (X, Y)) + grads.append((gX.detach().cpu().numpy(), gY.detach().cpu().numpy())) + + ref_gX, ref_gY = grads[0] + for gX, gY in grads[1:]: + np.testing.assert_allclose(gX, ref_gX, rtol=1e-5, atol=1e-8) + np.testing.assert_allclose(gY, ref_gY, rtol=1e-5, atol=1e-8) + + def test_longer_sequence_grad_is_finite(self): + rng = np.random.default_rng(4) + X = torch.tensor(rng.standard_normal((96, 2)), dtype=torch.float64, device=CPU, requires_grad=True) + Y = torch.tensor(rng.standard_normal((80, 2)), dtype=torch.float64, device=CPU, requires_grad=True) + out = compute_sig_kernel_fast_diff(self.ps, X, Y) + gX, gY = torch.autograd.grad(out, (X, Y)) + self.assertTrue(torch.isfinite(gX).all()) + self.assertTrue(torch.isfinite(gY).all()) + + def test_prefix_family_fast_diff_alias_matches_prefix_family(self): + rng = np.random.default_rng(5) + X = torch.tensor(rng.standard_normal((7, 2)), dtype=torch.float64, device=CPU, requires_grad=True) + refs = torch.tensor(rng.standard_normal((2, 6, 2)), dtype=torch.float64, device=CPU, requires_grad=True) + G = torch.tensor(rng.standard_normal((6, 2)), dtype=torch.float64, device=CPU) + + out_naive = compute_prefix_family(self.ps, X, refs, min_prefix_len=2, max_prefix_len=7) + out_fast = compute_prefix_family_fast_diff(self.ps, X, refs, min_prefix_len=2, max_prefix_len=7) + np.testing.assert_allclose(out_naive.detach().cpu().numpy(), out_fast.detach().cpu().numpy(), rtol=1e-10, atol=1e-12) + + loss_naive = torch.sum(G * out_naive) + gX_naive, gRefs_naive = torch.autograd.grad(loss_naive, (X, refs), retain_graph=True) + loss_fast = torch.sum(G * out_fast) + gX_fast, gRefs_fast = torch.autograd.grad(loss_fast, (X, refs)) + + np.testing.assert_allclose( + gX_fast.detach().cpu().numpy(), + gX_naive.detach().cpu().numpy(), + rtol=1e-5, + atol=1e-8, + ) + np.testing.assert_allclose( + gRefs_fast.detach().cpu().numpy(), + gRefs_naive.detach().cpu().numpy(), + rtol=1e-5, + atol=1e-8, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_core_torch.py b/tests/test_core_torch.py new file mode 100644 index 0000000..a6a04c8 --- /dev/null +++ b/tests/test_core_torch.py @@ -0,0 +1,286 @@ +"""CPU-compatible tests for the core Torch signature-kernel implementation.""" + +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +import torch +from jax import value_and_grad + +from powersig.jax import static_kernels as jax_static_kernels +from powersig.jax.algorithm import PowerSigJax +from powersig.torch import static_kernels as torch_static_kernels +from powersig.torch.algorithm import ( + PowerSigTorch, + _round_to_power_of_2, + build_stencil, + compute_block_size, + compute_vandermonde_vectors, + estimate_bytes_per_pair, + get_available_gpu_memory, + get_diagonal_range_bool, +) +from powersig.util.grid import get_diagonal_range as get_diagonal_range_py + + +CPU = torch.device("cpu") +JAX_CPU = jax.devices("cpu")[0] + + +class TestBuildStencil(unittest.TestCase): + def setUp(self): + self.order = 4 + self.dtype = torch.float64 + + def test_shape(self): + stencil = build_stencil(self.order, self.dtype) + self.assertEqual(stencil.shape, (self.order, self.order)) + + def test_first_row_and_column_are_ones(self): + stencil = build_stencil(self.order, self.dtype) + np.testing.assert_allclose(stencil[0].cpu().numpy(), np.ones(self.order)) + np.testing.assert_allclose(stencil[:, 0].cpu().numpy(), np.ones(self.order)) + + def test_known_values(self): + stencil = build_stencil(self.order, self.dtype) + expected = np.array( + [ + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 0.5, 1.0 / 3.0], + [1.0, 0.5, 0.25, 1.0 / 12.0], + [1.0, 1.0 / 3.0, 1.0 / 12.0, 1.0 / 36.0], + ] + ) + np.testing.assert_allclose(stencil.cpu().numpy(), expected, rtol=1e-10) + + +class TestVandermondeVectors(unittest.TestCase): + def test_unit_step(self): + v_s, v_t = compute_vandermonde_vectors(1.0, 1.0, 4, torch.float64) + np.testing.assert_allclose(v_s.cpu().numpy(), np.ones(4)) + np.testing.assert_allclose(v_t.cpu().numpy(), np.ones(4)) + + def test_power_scaling(self): + v_s, v_t = compute_vandermonde_vectors(0.5, 0.25, 4, torch.float64) + np.testing.assert_allclose(v_s.cpu().numpy(), [1.0, 0.5, 0.25, 0.125], rtol=1e-10) + np.testing.assert_allclose( + v_t.cpu().numpy(), [1.0, 0.25, 0.0625, 0.015625], rtol=1e-10 + ) + + +class TestDiagonalRange(unittest.TestCase): + def test_boolean_geometry_matches_python_reference(self): + for rows, cols in [(3, 2), (2, 4), (5, 5)]: + for d in range(rows + cols - 1): + expected = get_diagonal_range_py(d, rows, cols) + self.assertEqual(get_diagonal_range_bool(d, rows, cols), expected) + + +class TestBlockSizeUtils(unittest.TestCase): + def test_round_to_power_of_2(self): + self.assertEqual(_round_to_power_of_2(1), 1) + self.assertEqual(_round_to_power_of_2(2), 2) + self.assertEqual(_round_to_power_of_2(3), 4) + self.assertEqual(_round_to_power_of_2(5), 8) + self.assertEqual(_round_to_power_of_2(16), 16) + self.assertEqual(_round_to_power_of_2(17), 32) + + def test_estimate_bytes_per_pair(self): + bpp = estimate_bytes_per_pair(100, 32, torch.float64) + self.assertEqual(bpp, 2_636_800) + + def test_compute_block_size_bounded(self): + bs = compute_block_size(100, 32, torch.float64, CPU, 1000) + self.assertGreaterEqual(bs, 1) + self.assertLessEqual(bs, 256) + self.assertEqual(bs & (bs - 1), 0) + + def test_cpu_memory_budget_is_positive(self): + self.assertGreater(get_available_gpu_memory(CPU), 0) + + +class TestGramMatrix(unittest.TestCase): + def setUp(self): + rng = np.random.default_rng(42) + self.ps = PowerSigTorch(order=8, device=CPU, dtype=torch.float64) + self.ps_jax = PowerSigJax(order=8, device=JAX_CPU, dtype=jnp.float64) + self.X = torch.tensor(rng.normal(size=(4, 10, 3)), dtype=torch.float64, device=CPU) + self.Y = torch.tensor(rng.normal(size=(4, 10, 3)), dtype=torch.float64, device=CPU) + + def test_block_size_1_matches_auto(self): + gram_seq = self.ps.compute_gram_matrix(self.X, self.Y, block_size=1, show_progress=False) + gram_auto = self.ps.compute_gram_matrix(self.X, self.Y, show_progress=False) + np.testing.assert_allclose(gram_seq.cpu().numpy(), gram_auto.cpu().numpy(), rtol=1e-10) + + def test_explicit_block_sizes_match(self): + gram_1 = self.ps.compute_gram_matrix(self.X, self.Y, block_size=1, show_progress=False) + gram_4 = self.ps.compute_gram_matrix(self.X, self.Y, block_size=4, show_progress=False) + gram_16 = self.ps.compute_gram_matrix(self.X, self.Y, block_size=16, show_progress=False) + np.testing.assert_allclose(gram_1.cpu().numpy(), gram_4.cpu().numpy(), rtol=1e-10) + np.testing.assert_allclose(gram_1.cpu().numpy(), gram_16.cpu().numpy(), rtol=1e-10) + + def test_symmetric(self): + gram = self.ps.compute_gram_matrix(self.X, self.X, symmetric=True, show_progress=False) + np.testing.assert_allclose(gram.cpu().numpy(), gram.t().cpu().numpy(), rtol=1e-10) + + def test_symmetric_matches_full(self): + gram_full = self.ps.compute_gram_matrix(self.X, self.X, symmetric=False, show_progress=False) + gram_sym = self.ps.compute_gram_matrix(self.X, self.X, symmetric=True, show_progress=False) + np.testing.assert_allclose(gram_full.cpu().numpy(), gram_sym.cpu().numpy(), rtol=1e-10) + + def test_single_entry_matches_gram(self): + gram = self.ps.compute_gram_matrix(self.X, self.Y, show_progress=False) + for i in range(2): + for j in range(2): + single = self.ps.compute_signature_kernel(self.X[i], self.Y[j]) + np.testing.assert_allclose( + float(gram[i, j]), float(single), rtol=1e-6, err_msg=f"Mismatch at ({i}, {j})" + ) + + def test_call_interface(self): + gram_method = self.ps.compute_gram_matrix(self.X, self.Y, show_progress=False) + gram_call = self.ps(self.X, self.Y, show_progress=False) + np.testing.assert_allclose(gram_method.cpu().numpy(), gram_call.cpu().numpy(), rtol=1e-10) + + def test_call_with_block_size(self): + gram = self.ps(self.X, self.Y, block_size=2, show_progress=False) + gram_ref = self.ps(self.X, self.Y, block_size=1, show_progress=False) + np.testing.assert_allclose(gram.cpu().numpy(), gram_ref.cpu().numpy(), rtol=1e-10) + + def test_boolean_and_minimum_forward_geometry_match(self): + min_val = self.ps.compute_signature_kernel(self.X[0], self.Y[0]) + bool_val = self.ps.compute_signature_kernel_bool_geometry(self.X[0], self.Y[0]) + np.testing.assert_allclose(float(min_val), float(bool_val), rtol=1e-10, atol=1e-12) + + def test_chunked_matches_forward(self): + direct = self.ps.compute_signature_kernel(self.X[0], self.Y[0]) + chunked = self.ps.compute_signature_kernel_chunked(self.X[0], self.Y[0]) + np.testing.assert_allclose(float(direct), float(chunked), rtol=1e-10, atol=1e-12) + + def test_asymmetric_lengths_transpose_symmetry_small(self): + rng = np.random.default_rng(7) + X = torch.tensor(rng.normal(size=(1, 4, 2)), dtype=torch.float64, device=CPU) + Y = torch.tensor(rng.normal(size=(1, 3, 2)), dtype=torch.float64, device=CPU) + xy = self.ps.compute_gram_matrix(X, Y, show_progress=False) + yx = self.ps.compute_gram_matrix(Y, X, show_progress=False) + np.testing.assert_allclose(xy.cpu().numpy(), yx.t().cpu().numpy(), rtol=1e-10, atol=1e-12) + + def test_asymmetric_lengths_transpose_symmetry_longer(self): + rng = np.random.default_rng(17) + X = torch.tensor(0.1 * rng.normal(size=(1, 80, 2)), dtype=torch.float64, device=CPU) + Y = torch.tensor(0.1 * rng.normal(size=(1, 70, 2)), dtype=torch.float64, device=CPU) + xy = self.ps.compute_gram_matrix(X, Y, show_progress=False) + yx = self.ps.compute_gram_matrix(Y, X, show_progress=False) + np.testing.assert_allclose(xy.cpu().numpy(), yx.t().cpu().numpy(), rtol=1e-10, atol=1e-12) + + def test_linear_kernel_matches_jax_reference(self): + X_np = np.asarray(self.X[:2].cpu().numpy()) + Y_np = np.asarray(self.Y[:2].cpu().numpy()) + torch_gram = self.ps.compute_gram_matrix(X_np, Y_np, show_progress=False).cpu().numpy() + jax_gram = np.asarray( + self.ps_jax.compute_gram_matrix( + jnp.asarray(X_np, dtype=jnp.float64), + jnp.asarray(Y_np, dtype=jnp.float64), + ) + ) + np.testing.assert_allclose(torch_gram, jax_gram, rtol=1e-9, atol=1e-11) + + def test_rbf_kernel_matches_jax_reference(self): + rng = np.random.default_rng(99) + X_np = rng.normal(size=(2, 8, 2)) + Y_np = rng.normal(size=(2, 7, 2)) + + torch_ps = PowerSigTorch( + order=6, static_kernel=torch_static_kernels.rbf_kernel, device=CPU, dtype=torch.float64 + ) + jax_ps = PowerSigJax( + order=6, static_kernel=jax_static_kernels.rbf_kernel, device=JAX_CPU, dtype=jnp.float64 + ) + + torch_gram = torch_ps.compute_gram_matrix(X_np, Y_np, show_progress=False).cpu().numpy() + jax_gram = np.asarray( + jax_ps.compute_gram_matrix( + jnp.asarray(X_np, dtype=jnp.float64), + jnp.asarray(Y_np, dtype=jnp.float64), + ) + ) + np.testing.assert_allclose(torch_gram, jax_gram, rtol=1e-8, atol=1e-10) + + +class TestAutograd(unittest.TestCase): + def test_linear_gradient_matches_jax(self): + rng = np.random.default_rng(123) + X_np = rng.normal(size=(8, 2)) + Y_np = rng.normal(size=(7, 2)) + + torch_ps = PowerSigTorch(order=4, device=CPU, dtype=torch.float64) + X_t = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + Y_t = torch.tensor(Y_np, dtype=torch.float64, device=CPU, requires_grad=True) + torch_val = torch_ps.compute_signature_kernel(X_t, Y_t) + torch_grad_x, torch_grad_y = torch.autograd.grad(torch_val, (X_t, Y_t)) + + jax_ps = PowerSigJax(order=4, device=JAX_CPU, dtype=jnp.float64) + + def jax_fn(x, y): + return jax_ps.compute_signature_kernel(x, y) + + jax_val, (jax_grad_x, jax_grad_y) = value_and_grad(jax_fn, argnums=(0, 1))( + jnp.asarray(X_np, dtype=jnp.float64), jnp.asarray(Y_np, dtype=jnp.float64) + ) + + np.testing.assert_allclose(float(torch_val.detach().cpu()), float(jax_val), rtol=1e-10, atol=1e-12) + np.testing.assert_allclose( + torch_grad_x.detach().cpu().numpy(), + np.asarray(jax_grad_x), + rtol=1e-5, + atol=1e-8, + ) + np.testing.assert_allclose( + torch_grad_y.detach().cpu().numpy(), + np.asarray(jax_grad_y), + rtol=1e-5, + atol=1e-8, + ) + + def test_rbf_gradient_matches_jax(self): + rng = np.random.default_rng(321) + X_np = 0.2 * rng.normal(size=(6, 2)) + Y_np = 0.2 * rng.normal(size=(5, 2)) + + torch_ps = PowerSigTorch( + order=4, static_kernel=torch_static_kernels.rbf_kernel, device=CPU, dtype=torch.float64 + ) + X_t = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + Y_t = torch.tensor(Y_np, dtype=torch.float64, device=CPU, requires_grad=True) + torch_val = torch_ps.compute_signature_kernel(X_t, Y_t) + torch_grad_x, torch_grad_y = torch.autograd.grad(torch_val, (X_t, Y_t)) + + jax_ps = PowerSigJax( + order=4, static_kernel=jax_static_kernels.rbf_kernel, device=JAX_CPU, dtype=jnp.float64 + ) + + def jax_fn(x, y): + return jax_ps.compute_signature_kernel(x, y) + + jax_val, (jax_grad_x, jax_grad_y) = value_and_grad(jax_fn, argnums=(0, 1))( + jnp.asarray(X_np, dtype=jnp.float64), jnp.asarray(Y_np, dtype=jnp.float64) + ) + + np.testing.assert_allclose(float(torch_val.detach().cpu()), float(jax_val), rtol=1e-9, atol=1e-11) + np.testing.assert_allclose( + torch_grad_x.detach().cpu().numpy(), + np.asarray(jax_grad_x), + rtol=2e-5, + atol=1e-8, + ) + np.testing.assert_allclose( + torch_grad_y.detach().cpu().numpy(), + np.asarray(jax_grad_y), + rtol=2e-5, + atol=1e-8, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_prefix_family_torch.py b/tests/test_prefix_family_torch.py new file mode 100644 index 0000000..7ce1b7f --- /dev/null +++ b/tests/test_prefix_family_torch.py @@ -0,0 +1,184 @@ +"""Tests for the Torch prefix-family primitive and its custom backward.""" + +import unittest +from functools import partial + +import numpy as np +import torch + +from powersig.torch import compute_prefix_family +from powersig.torch.algorithm import PowerSigTorch +from powersig.torch.static_kernels import linear_kernel, rbf_kernel + + +CPU = torch.device("cpu") + + +def _scalar_ref(ps, X, Y): + return float(ps.compute_gram_matrix(X[None], Y[None], show_progress=False)[0, 0]) + + +class TestPrefixFamilyForward(unittest.TestCase): + def setUp(self): + np.random.seed(0) + + def _check(self, ps, Tx, Ty, R, d, *, tol=1e-10): + rng = np.random.default_rng(0) + X = torch.tensor(rng.standard_normal((Tx, d)), dtype=torch.float64, device=CPU) + refs = torch.tensor(rng.standard_normal((R, Ty, d)), dtype=torch.float64, device=CPU) + out = compute_prefix_family(ps, X, refs, min_prefix_len=2, max_prefix_len=Tx) + self.assertEqual(tuple(out.shape), (Tx - 1, R)) + for k in range(2, Tx + 1): + for r in range(R): + ref_val = _scalar_ref(ps, X[:k], refs[r]) + my_val = float(out[k - 2, r]) + rel = abs(ref_val - my_val) / max(1e-10, abs(ref_val)) + self.assertLess( + rel, + tol, + f"Tx={Tx}, Ty={Ty}, k={k}, r={r}: mine={my_val}, ref={ref_val}, rel_err={rel:.2e}", + ) + + def test_linear_rows_lt_cols(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + self._check(ps, Tx=4, Ty=6, R=3, d=3) + + def test_linear_rows_gt_cols(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + self._check(ps, Tx=7, Ty=4, R=3, d=3) + + def test_linear_square(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + self._check(ps, Tx=5, Ty=5, R=3, d=3) + + def test_rbf_asymmetric(self): + ps = PowerSigTorch( + order=5, + static_kernel=partial(rbf_kernel, bandwidth=0.75), + device=CPU, + dtype=torch.float64, + ) + self._check(ps, Tx=6, Ty=4, R=2, d=2, tol=1e-9) + + def test_single_reference_single_prefix_matches_scalar(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + rng = np.random.default_rng(7) + X = torch.tensor(rng.standard_normal((5, 3)), dtype=torch.float64, device=CPU) + Y = torch.tensor(rng.standard_normal((7, 3)), dtype=torch.float64, device=CPU) + out = compute_prefix_family(ps, X, Y[None], min_prefix_len=5, max_prefix_len=5) + self.assertEqual(tuple(out.shape), (1, 1)) + ref = _scalar_ref(ps, X, Y) + self.assertLess(abs(float(out[0, 0]) - ref) / max(1e-10, abs(ref)), 1e-10) + + def test_sub_range(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + rng = np.random.default_rng(2) + X = torch.tensor(rng.standard_normal((8, 3)), dtype=torch.float64, device=CPU) + refs = torch.tensor(rng.standard_normal((2, 10, 3)), dtype=torch.float64, device=CPU) + out = compute_prefix_family(ps, X, refs, min_prefix_len=4, max_prefix_len=7) + self.assertEqual(tuple(out.shape), (4, 2)) + for k in range(4, 8): + for r in range(2): + ref_val = _scalar_ref(ps, X[:k], refs[r]) + self.assertLess( + abs(float(out[k - 4, r]) - ref_val) / max(1e-10, abs(ref_val)), + 1e-10, + ) + + def test_min_prefix_len_below_2_raises(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + X = torch.zeros((4, 2), dtype=torch.float64, device=CPU) + refs = torch.zeros((1, 4, 2), dtype=torch.float64, device=CPU) + with self.assertRaises(ValueError): + compute_prefix_family(ps, X, refs, min_prefix_len=1) + + def test_max_prefix_len_exceeds_T_raises(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + X = torch.zeros((4, 2), dtype=torch.float64, device=CPU) + refs = torch.zeros((1, 4, 2), dtype=torch.float64, device=CPU) + with self.assertRaises(ValueError): + compute_prefix_family(ps, X, refs, max_prefix_len=5) + + +class TestPrefixFamilyGrad(unittest.TestCase): + def _naive_prefix_family(self, ps, X, refs, min_prefix_len, max_prefix_len): + rows = [] + for k in range(min_prefix_len, max_prefix_len + 1): + row = [] + for r in range(refs.shape[0]): + row.append(ps.compute_signature_kernel(X[:k], refs[r])) + rows.append(torch.stack(row)) + return torch.stack(rows) + + def test_grad_wrt_state_path_matches_naive(self): + ps = PowerSigTorch(order=4, static_kernel=linear_kernel, device=CPU, dtype=torch.float64) + rng = np.random.default_rng(3) + X_np = rng.standard_normal((4, 2)) + refs_np = rng.standard_normal((2, 5, 2)) + G = torch.tensor(rng.standard_normal((3, 2)), dtype=torch.float64, device=CPU) + + X_custom = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + refs_custom = torch.tensor(refs_np, dtype=torch.float64, device=CPU, requires_grad=True) + out_custom = compute_prefix_family(ps, X_custom, refs_custom, min_prefix_len=2, max_prefix_len=4) + loss_custom = torch.sum(G * out_custom) + gX_custom, gRefs_custom = torch.autograd.grad(loss_custom, (X_custom, refs_custom)) + + X_naive = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + refs_naive = torch.tensor(refs_np, dtype=torch.float64, device=CPU, requires_grad=True) + out_naive = self._naive_prefix_family(ps, X_naive, refs_naive, 2, 4) + loss_naive = torch.sum(G * out_naive) + gX_naive, gRefs_naive = torch.autograd.grad(loss_naive, (X_naive, refs_naive)) + + np.testing.assert_allclose( + gX_custom.detach().cpu().numpy(), + gX_naive.detach().cpu().numpy(), + rtol=1e-5, + atol=1e-8, + ) + np.testing.assert_allclose( + gRefs_custom.detach().cpu().numpy(), + gRefs_naive.detach().cpu().numpy(), + rtol=1e-5, + atol=1e-8, + ) + + def test_grad_wrt_state_path_and_refs_matches_naive_rbf(self): + ps = PowerSigTorch( + order=4, + static_kernel=partial(rbf_kernel, bandwidth=0.8), + device=CPU, + dtype=torch.float64, + ) + rng = np.random.default_rng(4) + X_np = 0.2 * rng.standard_normal((3, 2)) + refs_np = 0.2 * rng.standard_normal((2, 4, 2)) + G = torch.tensor(rng.standard_normal((2, 2)), dtype=torch.float64, device=CPU) + + X_custom = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + refs_custom = torch.tensor(refs_np, dtype=torch.float64, device=CPU, requires_grad=True) + out_custom = compute_prefix_family(ps, X_custom, refs_custom, min_prefix_len=2, max_prefix_len=3) + loss_custom = torch.sum(G * out_custom) + gX_custom, gRefs_custom = torch.autograd.grad(loss_custom, (X_custom, refs_custom)) + + X_naive = torch.tensor(X_np, dtype=torch.float64, device=CPU, requires_grad=True) + refs_naive = torch.tensor(refs_np, dtype=torch.float64, device=CPU, requires_grad=True) + out_naive = self._naive_prefix_family(ps, X_naive, refs_naive, 2, 3) + loss_naive = torch.sum(G * out_naive) + gX_naive, gRefs_naive = torch.autograd.grad(loss_naive, (X_naive, refs_naive)) + + np.testing.assert_allclose( + gX_custom.detach().cpu().numpy(), + gX_naive.detach().cpu().numpy(), + rtol=2e-5, + atol=1e-8, + ) + np.testing.assert_allclose( + gRefs_custom.detach().cpu().numpy(), + gRefs_naive.detach().cpu().numpy(), + rtol=2e-5, + atol=1e-8, + ) + + +if __name__ == "__main__": + unittest.main() From 9b53296b2b8bf54fc5293088d6158b92c41b11ab Mon Sep 17 00:00:00 2001 From: Matthew Tamayo-Rios Date: Wed, 26 Aug 2026 03:06:53 -0700 Subject: [PATCH 3/5] Document JAX and PyTorch as selectable backends The README only showed the JAX path, so there was no way to tell that a PyTorch backend existed or how to install it. Install now leads with the per-backend extras, Getting Started carries a runnable quickstart for each, and a "Choosing a backend" section maps the two APIs onto each other and says when to reach for which. Adds examples/simple_torch.py as the PyTorch mirror of examples/simple.py. Install commands switched to the published PyPI package with extras; the previous plain git URL installed no backend at all, and the #egg= form does not carry extras. CuPy is described as forward-Gram only, which is what powersig.cupy_backend actually exposes. --- README.md | 153 ++++++++++++++++++++++++++++++++++++--- examples/simple_torch.py | 31 ++++++++ 2 files changed, 175 insertions(+), 9 deletions(-) create mode 100644 examples/simple_torch.py diff --git a/README.md b/README.md index d449cfc..d3d47c4 100644 --- a/README.md +++ b/README.md @@ -5,16 +5,56 @@ Using ADM-derived Neumann series to compute signature kernels. -## Installation +PowerSig ships two interchangeable backends — **JAX** and **PyTorch** — behind the +same API. Pick whichever matches the framework you already use; both produce the +same kernel values to machine precision (see [Choosing a backend](#choosing-a-backend)). + +## Installation + +Requires Python 3.12+. Install the extra for the backend you want — the backend +frameworks are optional dependencies, so nothing heavyweight is pulled in by default. + +```bash +# JAX, CPU only +pip install "powersig[jax-cpu]" + +# JAX, CUDA 13 GPU +pip install "powersig[jax-gpu]" + +# PyTorch (CPU or CUDA, depending on the torch wheel you install) +pip install "powersig[torch]" +``` + +To install from source, use the same extras with a direct reference: + ```bash -pip install git+https://github.com/geekbeast/powersig.git +pip install "powersig[torch] @ git+https://github.com/geekbeast/powersig.git" ``` -Requires Python 3.12+ +| Extra | Backend | Pulls in | +| --- | --- | --- | +| `jax-cpu` | JAX | `jax[cpu]>=0.4.34` | +| `jax-gpu` | JAX | `jax[cuda13]>=0.10.0` | +| `torch` | PyTorch | `torch>=2.5.0` | +| `cupy` / `cupy-cuda13` | CuPy | `cupy-cuda12x` / `cupy-cuda13x` | +| `all` | JAX + PyTorch + CuPy | all of the above | -Requires PyTorch 2.5+, JAX 0.6.0+, or cupy 13.4.1+ depending on which implementation you prefer. +For a specific PyTorch build (a CPU-only wheel, or a CUDA version other than the +PyPI default), install `torch` first from the PyTorch index and then install +PowerSig: + +```bash +pip install torch --index-url https://download.pytorch.org/whl/cpu +pip install "powersig[torch]" +``` + +Importing `powersig` does not import any backend, so having only one installed is +fine — backends are resolved lazily on first use. ## Getting Started + +### JAX + ```python import jax.numpy as jnp from powersig.jax.utils import fractional_brownian_motion @@ -25,7 +65,7 @@ def main(): n_steps = 1000 n_paths = 2 hurst = 0.7 - + # Generate fBM using the jax wrapper fbm_paths, dt = fractional_brownian_motion( n_steps=n_steps, @@ -33,13 +73,13 @@ def main(): hurst=hurst, dim=1 ) - + # Initialize PowerSigJax with polynomial order 8 powersig = PowerSigJax(order=8) - + # Compute the signature kernel kernel_matrix = powersig(fbm_paths) - + print("Shape of fBM paths:", fbm_paths.shape) print("Shape of kernel matrix:", kernel_matrix.shape) print("\nKernel matrix:") @@ -48,4 +88,99 @@ def main(): if __name__ == "__main__": main() ``` -This example is also available in the repo under [examples](examples/simple.py). \ No newline at end of file + +This example is also available in the repo under [examples/simple.py](examples/simple.py). + +### PyTorch + +```python +import torch +from powersig.torch.utils import fractional_brownian_motion +from powersig.torch.algorithm import PowerSigTorch + +def main(): + # Generate fBM paths + n_steps = 1000 + n_paths = 2 + hurst = 0.7 + + # Generate fBM using the torch wrapper + fbm_paths, dt = fractional_brownian_motion( + n_steps=n_steps, + n_paths=n_paths, + hurst=hurst, + dim=1 + ) + + # Initialize PowerSigTorch with polynomial order 8 + powersig = PowerSigTorch(order=8) + + # Compute the signature kernel + kernel_matrix = powersig(fbm_paths) + + print("Shape of fBM paths:", tuple(fbm_paths.shape)) + print("Shape of kernel matrix:", tuple(kernel_matrix.shape)) + print("\nKernel matrix:") + print(kernel_matrix) + +if __name__ == "__main__": + main() +``` + +This example is also available in the repo under [examples/simple_torch.py](examples/simple_torch.py). + +## Choosing a backend + +The two backends expose the same surface, so switching is a matter of swapping the +import and the array type: + +| | JAX | PyTorch | +| --- | --- | --- | +| Estimator | `powersig.jax.algorithm.PowerSigJax` | `powersig.torch.algorithm.PowerSigTorch` | +| fBM helper | `powersig.jax.utils.fractional_brownian_motion` | `powersig.torch.utils.fractional_brownian_motion` | +| Static kernels | `powersig.jax.static_kernels` | `powersig.torch.static_kernels` | +| Gram matrix | `ps(X)` / `ps(X, Y)` | `ps(X)` / `ps(X, Y)` | +| Single pair | `ps.compute_signature_kernel(x, y)` | `ps.compute_signature_kernel(x, y)` | +| Gradients | `jax.grad` / `jax.value_and_grad` | `torch.autograd` (`.backward()`) | + +Both constructors take the same arguments: + +```python +PowerSigJax(order=32, static_kernel=linear_kernel, device=None, dtype=jnp.float64) +PowerSigTorch(order=32, static_kernel=linear_kernel, device=None, dtype=torch.float64) +``` + +- `order` — truncation order of the power series. Higher is more accurate and more + expensive; 8–32 is the usual range. +- `static_kernel` — the static kernel lifted to a signature kernel. `linear_kernel` + (default) and `rbf_kernel` are provided by each backend's `static_kernels` module. +- `device` — defaults to the first available GPU, else CPU. +- `dtype` — defaults to float64. Use float32 to trade accuracy for speed. + +Paths are `(batch, length, dim)` for Gram matrices and `(length, dim)` for a single +pair. The two paths in a pair need not have the same length. + +Which one to pick: + +- **PyTorch** if your model, data loading, or training loop is already in PyTorch — + the kernel is differentiable through `torch.autograd`, so it drops into an + existing training loop without a framework boundary. +- **JAX** if you want `jit`/`vmap`/`grad` composition, or are already in a JAX + codebase. + +A **CuPy** backend also exists under `powersig.cupy_backend`. It covers the forward +Gram computation only — no autodiff and no pluggable static kernel — so the JAX and +PyTorch backends are the supported choices for general use. + +## Testing + +```bash +pip install ".[jax-cpu,dev]" # add "torch" for the PyTorch suite +pytest tests/test_core_jax.py # JAX backend +pytest tests/test_core_torch.py \ + tests/test_autodiff_torch.py \ + tests/test_prefix_family_torch.py # PyTorch backend +``` + +The PyTorch suite cross-checks its results against the JAX implementation, so it +needs both backends installed. CI runs both on every push and pull request. diff --git a/examples/simple_torch.py b/examples/simple_torch.py new file mode 100644 index 0000000..76442a2 --- /dev/null +++ b/examples/simple_torch.py @@ -0,0 +1,31 @@ +import torch +from powersig.torch.utils import fractional_brownian_motion +from powersig.torch.algorithm import PowerSigTorch + +def main(): + # Generate fBM paths + n_steps = 1000 + n_paths = 2 + hurst = 0.7 + + # Generate fBM using the torch wrapper + fbm_paths, dt = fractional_brownian_motion( + n_steps=n_steps, + n_paths=n_paths, + hurst=hurst, + dim=1 + ) + + # Initialize PowerSigTorch with polynomial order 8 + powersig = PowerSigTorch(order=8) + + # Compute the signature kernel + kernel_matrix = powersig(fbm_paths) + + print("Shape of fBM paths:", tuple(fbm_paths.shape)) + print("Shape of kernel matrix:", tuple(kernel_matrix.shape)) + print("\nKernel matrix:") + print(kernel_matrix) + +if __name__ == "__main__": + main() From f56abe7fe45ddf87025671800f3705d936879358 Mon Sep 17 00:00:00 2001 From: Matthew Tamayo-Rios Date: Wed, 26 Aug 2026 12:14:45 -0700 Subject: [PATCH 4/5] Guard that the corrected sweep geometry stays free The sweep runs this arithmetic once per anti-diagonal inside the hot loop, so a correctness fix that reached for jnp.where, a Python-level branch, or a helper call would pay for itself there. The geometry correction on feature/custom-autodiff took that route -- it swapped the inline arithmetic for a get_diagonal_range call and added a boolean-arithmetic variant plus timing coverage to compare the two. This fix keeps the inline boolean form and only corrects the indices, so the operation count is unchanged. Measured on the compiled module for a 129-point pair at order 8, main and this branch come out at 245819 and 245811 flops; the 8-flop delta is the tracer-leak fix, and isolating the geometry change alone reproduces main's 245819 exactly. The test compares the corrected expression against the one it replaced, compiling both in-process with the same JAX build rather than asserting a recorded number, so it does not drift as JAX changes its cost model. A second test pins the expression to powersig/util/grid.py so the copy in the test cannot silently diverge from the shipped geometry. --- tests/test_core_jax.py | 76 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/tests/test_core_jax.py b/tests/test_core_jax.py index 1209f21..2a64dae 100644 --- a/tests/test_core_jax.py +++ b/tests/test_core_jax.py @@ -310,5 +310,81 @@ def upsample(P, factor): np.testing.assert_allclose(asymmetric, reference, rtol=1e-10, atol=1e-12) +# --------------------------------------------------------------------------- +# Cost of the corrected sweep geometry +# --------------------------------------------------------------------------- +class TestGeometryCost(unittest.TestCase): + """The asymmetric-geometry fix must stay free. + + The sweep runs this arithmetic once per anti-diagonal inside the hot loop, so + a correctness fix that reached for jnp.where, a Python-level branch, or a + helper call could pay for itself in the inner loop. Measured on the compiled + module, the shipped correction costs exactly the same as the original + expression it replaced. + + Both are compiled in-process with the same JAX build and compared to each + other rather than to a recorded number, so the assertion does not drift as + JAX changes its cost model. + """ + + ROWS, COLS, N_DIAGONALS = 128, 96, 256 + + @staticmethod + def _original_geometry(d, rows, cols): + # The expression this fix replaced: correct only when rows == cols, but + # it is the cheapness baseline the corrected version has to match. + t_start = (d < cols) * 0 + (d >= cols) * (d - cols + 1) + s_start = (d < cols) * d + (d >= cols) * (cols - 1) + return s_start, t_start, jnp.minimum(rows - t_start, s_start + 1) + + @staticmethod + def _corrected_geometry(d, rows, cols): + # Must stay in step with powersig/jax/algorithm.py::compute_diagonal and + # powersig/util/grid.py; test_corrected_geometry_matches_shared_helper + # fails if it drifts. + s_start = (d < rows) * d + (d >= rows) * (rows - 1) + t_start = (d < rows) * 0 + (d >= rows) * (d - rows + 1) + return s_start, t_start, jnp.minimum(s_start + 1, cols - t_start) + + def _compiled_cost(self, fn): + ds = jnp.arange(self.N_DIAGONALS, dtype=jnp.int32) + compiled = jax.jit( + lambda d: jax.vmap(lambda x: fn(x, self.ROWS, self.COLS))(d) + ).lower(ds).compile() + analysis = compiled.cost_analysis() + if isinstance(analysis, list): + analysis = analysis[0] + return analysis.get("flops"), analysis.get("bytes accessed") + + def test_corrected_geometry_matches_shared_helper(self): + """Guards the copy above against drifting from the shipped geometry.""" + for rows in range(1, 7): + for cols in range(1, 7): + for d in range(rows + cols - 1): + got = tuple( + int(v) for v in self._corrected_geometry(d, rows, cols) + ) + self.assertEqual( + got, + get_diagonal_range(d, rows, cols), + msg=f"d={d} rows={rows} cols={cols}", + ) + + def test_corrected_geometry_costs_no_more_than_original(self): + original_flops, original_bytes = self._compiled_cost(self._original_geometry) + corrected_flops, corrected_bytes = self._compiled_cost(self._corrected_geometry) + + self.assertIsNotNone(original_flops) + self.assertLessEqual( + corrected_flops, + original_flops, + msg=( + f"corrected geometry costs {corrected_flops} flops vs " + f"{original_flops} for the expression it replaced" + ), + ) + self.assertLessEqual(corrected_bytes, original_bytes) + + if __name__ == "__main__": unittest.main() From b7857e483201dbb61c211fd554a2d4aab4d3f318 Mon Sep 17 00:00:00 2001 From: Matthew Tamayo-Rios Date: Wed, 26 Aug 2026 13:00:14 -0700 Subject: [PATCH 5/5] Stop compile_forward from aliasing the CUDA-graph output buffer compile_forward=True routes the sweep through torch.compile with mode="reduce-overhead", which replays a CUDA graph writing into a fixed output buffer. compute_signature_kernel returned a view of that buffer, so a result the caller was still holding got silently overwritten by their next call: ps = PowerSigTorch(order=8, compile_forward=True) ks = [ps.compute_signature_kernel(x, y) for x, y in pairs] # every entry of ks held the LAST pair's value Values were correct when consumed immediately, which is why timing code never trips over this -- a benchmark reads the number straight away. Measured drift on held tensors was 1.49 absolute before this change and 0.0 after. compute_gram_matrix was never affected: it consumes each entry into the output matrix before the next call. Confirmed at 2.7e-15 against the eager path, symmetric, every entry matching an independent pairwise computation. The fix clones the scalar out of the graph buffer before returning it. The output is a scalar, so the copy costs nothing next to the sweep, and the CPU path is untouched since compile_forward only applies on CUDA. Also documents the flag in the README, because leaving it off is expensive: eager, the sweep costs a flat ~257us per anti-diagonal whatever the truncation order -- it is bound by kernel-launch overhead, not arithmetic, so order 8 and order 32 run at the same speed. On an RTX 4090 at 129 points, order 8, that is 65.7ms eager vs 5.6ms compiled vs 6.4ms for JAX. The flag is not free to enable (CUDA only, and dynamic=False means a compile pause of tens of seconds per input shape), which is presumably why it defaults to off. --- README.md | 33 ++++++++++++++++++++++ powersig/torch/algorithm.py | 17 ++++++++++- tests/test_core_torch.py | 56 +++++++++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index d3d47c4..1cb5fb4 100644 --- a/README.md +++ b/README.md @@ -168,6 +168,39 @@ Which one to pick: - **JAX** if you want `jit`/`vmap`/`grad` composition, or are already in a JAX codebase. +Both compute the same values, so this is a question of which framework you are +already in rather than which is faster — with one caveat below. + +### PyTorch on a GPU: pass `compile_forward=True` + +```python +ps = PowerSigTorch(order=8, compile_forward=True) +``` + +One kernel between two 129-point 2-D paths at order 8, single RTX 4090, median +of repeated trials: + +| backend | time | +| --- | --- | +| PyTorch, `compile_forward=True` | 5.6 ms | +| JAX | 6.4 ms | +| PyTorch, default | 65.7 ms | + +Left alone, the PyTorch sweep is bound by per-anti-diagonal kernel-launch +overhead rather than by arithmetic — the cost sits at roughly 257 us per +anti-diagonal whatever the truncation order, so order 8 and order 32 run at the +same speed and the backend lands about 10x behind JAX. `compile_forward=True` +routes the sweep through `torch.compile` and recovers roughly 12x, putting it +level with or slightly ahead of JAX. + +It is off by default because it is not free to turn on: it applies only on CUDA, +and it compiles per input shape, costing a pause of tens of seconds on the first +call for each new path length. That is worth it for repeated work at a fixed +size, and not worth it for a handful of one-off kernels at varying lengths. + +Gram matrices amortize the launch overhead across pairs, so the default setting +is far less punishing there than it is for single pairs. + A **CuPy** backend also exists under `powersig.cupy_backend`. It covers the forward Gram computation only — no autodiff and no pluggable static kernel — so the JAX and PyTorch backends are the supported choices for general use. diff --git a/powersig/torch/algorithm.py b/powersig/torch/algorithm.py index fcc8d61..bce7b15 100644 --- a/powersig/torch/algorithm.py +++ b/powersig/torch/algorithm.py @@ -248,7 +248,22 @@ def compute_signature_kernel( ) -> torch.Tensor: X_i = self._as_single_path(X, device=device) Y_j = self._as_single_path(Y, device=device) - return self._minimum_sweep(X_i[None, ...], Y_j[None, ...])[0] + return self._detach_from_graph_buffer( + self._minimum_sweep(X_i[None, ...], Y_j[None, ...])[0] + ) + + def _detach_from_graph_buffer(self, out: torch.Tensor) -> torch.Tensor: + """Copy a compiled sweep's result out of its CUDA-graph static buffer. + + compile_forward uses torch.compile(mode="reduce-overhead"), which replays + a CUDA graph writing into a fixed output buffer. Without this copy the + tensor a caller is holding is silently overwritten by their next call, so + collecting results in a list yields the last value repeated. The clone is + a scalar, so it costs nothing next to the sweep itself. + """ + if self.compile_forward and out.is_cuda: + return out.clone() + return out def compute_signature_kernel_bool_geometry( self, X: torch.Tensor, Y: torch.Tensor, device: Optional[torch.device] = None diff --git a/tests/test_core_torch.py b/tests/test_core_torch.py index a6a04c8..120c93d 100644 --- a/tests/test_core_torch.py +++ b/tests/test_core_torch.py @@ -282,5 +282,61 @@ def jax_fn(x, y): ) +@unittest.skipUnless(torch.cuda.is_available(), "compile_forward only applies on CUDA") +class TestCompiledForward(unittest.TestCase): + """compile_forward=True must be as correct as it is fast. + + It routes the sweep through torch.compile(mode="reduce-overhead"), which + replays a CUDA graph writing into a fixed output buffer. The hazard is that + a result a caller is still holding gets overwritten by their next call. + """ + + N_PTS, ORDER = 65, 8 + + def setUp(self): + self.dev = torch.device("cuda") + self.eager = PowerSigTorch(order=self.ORDER, device=self.dev, + dtype=torch.float64, compile_forward=False) + self.compiled = PowerSigTorch(order=self.ORDER, device=self.dev, + dtype=torch.float64, compile_forward=True) + rng = np.random.default_rng(0) + self.pairs = [ + (torch.tensor(0.1 * rng.normal(size=(self.N_PTS, 2)).cumsum(0), device=self.dev), + torch.tensor(0.1 * rng.normal(size=(self.N_PTS, 2)).cumsum(0), device=self.dev)) + for _ in range(6) + ] + + def test_compiled_matches_eager(self): + for i, (X, Y) in enumerate(self.pairs): + np.testing.assert_allclose( + float(self.compiled.compute_signature_kernel(X, Y)), + float(self.eager.compute_signature_kernel(X, Y)), + rtol=1e-9, atol=1e-11, err_msg=f"pair {i}", + ) + + def test_held_results_survive_later_calls(self): + """Regression: results used to alias the CUDA-graph output buffer, so + collecting them in a list yielded the last value repeated.""" + held = [self.compiled.compute_signature_kernel(X, Y) for X, Y in self.pairs[:3]] + before = [float(h) for h in held] + for X, Y in self.pairs[3:]: + self.compiled.compute_signature_kernel(X, Y) + after = [float(h) for h in held] + self.assertEqual(before, after, "held results were overwritten by later calls") + + expected = [float(self.eager.compute_signature_kernel(X, Y)) + for X, Y in self.pairs[:3]] + np.testing.assert_allclose(after, expected, rtol=1e-9, atol=1e-11) + + def test_gram_matrix_matches_eager(self): + rng = np.random.default_rng(1) + X = torch.tensor(0.1 * rng.normal(size=(5, self.N_PTS, 2)).cumsum(1), device=self.dev) + np.testing.assert_allclose( + self.compiled.compute_gram_matrix(X, show_progress=False).cpu().numpy(), + self.eager.compute_gram_matrix(X, show_progress=False).cpu().numpy(), + rtol=1e-9, atol=1e-11, + ) + + if __name__ == "__main__": unittest.main()