Make the PyTorch backend work, and document JAX/PyTorch as selectable backends

.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

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,21 +65,21 @@ 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,
         n_paths=n_paths,
         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).
+
+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.

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()

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)*0 + (d>=cols)*(d-cols +1)
-            s_start = (d<cols)*d + (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)*d + (d>=rows)*(rows - 1)
+            t_start = (d<rows)*0 + (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)*0 + (diagonal_index>=cols)*(diagonal_index-cols +1)
-                s_start = (diagonal_index<cols)*diagonal_index + (diagonal_index>=cols)*(cols - 1)
+                s_start = (diagonal_index<rows)*diagonal_index + (diagonal_index>=rows)*(rows - 1)
+                t_start = (diagonal_index<rows)*0 + (diagonal_index>=rows)*(diagonal_index-rows + 1)
 
                 is_before_wrap = diagonal_index < rows
                 # rho = jax_compute_dot_prod_batch(jnp.take(dX_i, s_start-diagonal_indices, axis=0, fill_value=0), jnp.take(dY_j, t_start+diagonal_indices, axis=0, fill_value=0))
@@ -810,9 +816,9 @@ def process_column(c):
 @jit
 def get_diagonal_range(d: int, rows: int, cols: int) -> Tuple[int, int, int]:
     # d, s_start, t_start are 0 based indexes while rows/cols are shapes.
-    t_start = jnp.where(d<cols, 0, d-cols +1)
-    s_start = jnp.where(d<cols, d, cols - 1)
-    dlen = jnp.minimum(rows - t_start, s_start + 1)
+    s_start = jnp.where(d<rows, d, rows - 1)
+    t_start = jnp.where(d<rows, 0, d-rows + 1)
+    dlen = jnp.minimum(s_start + 1, cols - t_start)
     # if d < cols:
     #     # if d < cols, then we haven't hit the right edge of the grid
     #     t_start = 0