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

[geekbeast]

What this is

The ask was to check whether the PyTorch implementation works and, if so, document JAX and PyTorch as selectable backends. It did not work — on main it had never run — so this makes it work first, then documents both.

The PyTorch backend on main was never functional

• PowerSigTorch had no __call__, so it could not be used like PowerSigJax.
• compute_signature_kernel passed a dtype into build_stencil_s's device parameter and died in dynamo on torch.ones(..., device=torch.float64).
• compute_gram_matrix was a copy-paste of the JAX version: it called jnp.zeros (never imported) and .at[0].set(1) on a torch tensor.
• compute_gram_entry returned from inside its outer loop after the first chunk, with debug prints still in place.

None of this was caught because there were no PyTorch tests and CI ran only tests/test_core_jax.py.

This takes the implementation from feature/custom-autodiff, which reaches API parity with the JAX backend, plus its three test modules. Only powersig/torch/* is taken — the branch's JAX rewrite and its pre-CUDA-13 pyproject.toml are deliberately left behind.

Two real JAX bugs surfaced on the way

Porting the torch tests, which cross-check against JAX, turned up two pre-existing defects in the JAX backend. Both are fixed here, minimally.

1. The diagonal sweep was wrong whenever the two paths had different lengths. The geometry keyed s_start/t_start off cols, which walks s past the bottom edge of the grid when rows != cols. Against a brute-force enumeration, the old convention was wrong for 110 of the (d, rows, cols) triples with rows, cols <= 6 — every one with rows != cols.

The signature kernel depends on the paths, not on how finely they are sampled, so upsampling a piecewise-linear path along its own segments must leave the value unchanged. For a 6-point and a 5-point path regridded to a common 21 points:

	equal-length reference	asymmetric 6×5	error
before	0.7953811447	0.8221965290	2.7e-2
after	0.7953811447	0.7953811447	7.8e-16

The error did not shrink with truncation order, confirming geometry rather than truncation.

2. A jit tracer leak. compute_signature_kernel is jitted but assigned to self.exponents, storing a tracer on the instance. A compute_signature_kernel followed by a compute_gram_matrix on the same instance died with InvalidInputException.

Both come with regression tests. All five new JAX tests fail on the previous code and pass here.

Verification

The ported backend was checked independently of its own test suite:

Check	Result
Linear path vs closed form I₀(2)	matches to 1e-12
Gram vs JAX, linear + RBF static kernels	max diff 5e-15
CUDA vs CPU	4e-15
Reverse-mode gradients vs central differences	1.4e-9 relative
float32 and float64	both work
Full suite (JAX + PyTorch)	66 passed

Also in here

• CI now runs a [jax, torch] backend matrix, so neither implementation can rot silently again. The torch job installs the JAX CPU wheel too, since the PyTorch tests cross-check against the JAX reference.
• README leads with per-backend install extras, carries a runnable quickstart for each, and adds a "Choosing a backend" section mapping the two APIs onto each other. Install commands now use the published PyPI package with extras — the previous plain git URL installed no backend at all.
• examples/simple_torch.py, the PyTorch mirror of examples/simple.py. Both were run end to end.

Note

powersig/cupy_backend also consumes the corrected get_diagonal_range, so its asymmetric-input behaviour changes too — in the same direction. It has no test coverage and CuPy was not exercised here.