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/README.md b/README.md index d449cfc..1cb5fb4 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,132 @@ 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. + +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. + +## 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() 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(d0 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 +DIAGONAL_CHUNK_SIZE = 1024 +JIT_BOUNDARY_THRESHOLD = 64 +_MAX_BLOCK_SIZE_SMALL = 256 +_MAX_BLOCK_SIZE_LARGE = 16384 - """ - 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) - return gram_matrix +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) -@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)) - # In-place multiplication: diagonal * coefficient * vandermonde_slice - diagonal_view.mul_(rho_diag) +def get_available_gpu_memory(device: torch.device) -> int: + if device.type != "cuda": + return 4 * 1024**3 - return U + try: + free, _ = torch.cuda.mem_get_info(_device_index(device)) + return int(free) + except Exception: + pass + + try: + props = torch.cuda.get_device_properties(_device_index(device)) + return int(props.total_memory) + except Exception: + return 8 * 1024**3 + + +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._detach_from_graph_buffer( + self._minimum_sweep(X_i[None, ...], Y_j[None, ...])[0] + ) -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)], + 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 + ) -> 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 ) - u = batch_ADM_for_diagonal(rho, u_buf, S, T, stencil) + 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 - 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_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, + ) + + 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, + ) + + 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/powersig/util/grid.py b/powersig/util/grid.py index 36a2bd4..cf8b12c 100644 --- a/powersig/util/grid.py +++ b/powersig/util/grid.py @@ -2,14 +2,13 @@ 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. - - 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_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_jax.py b/tests/test_core_jax.py index 4b5d7d7..2a64dae 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,135 @@ 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) + + +# --------------------------------------------------------------------------- +# 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() diff --git a/tests/test_core_torch.py b/tests/test_core_torch.py new file mode 100644 index 0000000..120c93d --- /dev/null +++ b/tests/test_core_torch.py @@ -0,0 +1,342 @@ +"""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, + ) + + +@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() 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()