Where things stand
The previous update covered making DistanceMatrix itself able to hold a GPU-resident buffer. That's the foundation; it doesn't do anything on its own unless the functions that consume a distance matrix know what to do with one. This update covers the first two: PERMANOVA and Mantel, opened as PR #2503 against the project's development branch.
The goal was narrow on purpose: keep the exact semantics both functions already have, and only change where the computation happens based on where the input data lives.
Dispatch, unchanged semantics
Both functions already dispatched on the engine= parameter for CPU execution (Cython or Numba). The new dispatch sits above that: a DistanceMatrix backed by NumPy keeps using the existing engine path untouched; a DistanceMatrix backed by a GPU-resident buffer (PyTorch, JAX, or CuPy) is computed through the Array API instead, so the matrix never has to leave its device.
DistanceMatrix only
Both functions keep their existing input contract: a DistanceMatrix, nothing else. A bare GPU array is not a valid input, matching the semantics both functions had before this work.
PERMANOVA
The within-group sum of squares and pseudo-F statistic reduce to a small number of Array API reductions, reusing the same permutation scheme as the reference implementation so p-values match exactly for a given seed.
Mantel
Pearson and Spearman both have an Array API compute path. Kendall's tau has no such path (SciPy's implementation is host-only), so that variant materializes a GPU-resident matrix to the host and proceeds as before.
The whole diff is additive: nothing in the existing NumPy/Cython/Numba code path was touched, only new branches were added ahead of it.
Fixing a limitation instead of shipping it
An early version of this work carried a real limitation: if the two distance matrices passed to mantel had their IDs in a different order, the Array API path raised, instead of reordering them the way the NumPy path already does. That was true when the code was first written, before the DistanceMatrix work landed.
Before writing that limitation into the PR description, it was worth checking whether it was still true. It wasn't. The DistanceMatrix work had already made the underlying reorder/filter utilities Array-API aware, which meant the existing id-alignment helper (_order_dms) worked on a GPU-resident matrix without any further changes — it just hadn't been wired back in yet. Wiring it in removed the limitation entirely: id reordering, the strict-subset case, and id lookup/remapping all now behave identically to the NumPy path, verified against a NumPy reference at floating-point precision.
The smaller lesson generalizes: a documented limitation is a claim, and claims are worth re-checking against the current code before they go into a PR description, not just carried forward from when they were first true.
Chasing exact parity
An independent review of the diff surfaced two gaps between the Array API path and the existing NumPy path, both since fixed:
- Near-constant input warning. The NumPy Mantel path warns when one of the two distance vectors is close enough to constant that the Pearson/Spearman correlation becomes numerically unstable. The Array API path computed the same correlation but skipped the warning. Fixed by adding the same threshold check, verified to fire identically on both paths for near-constant input.
- Untested tie-handling. The existing test suite for the Array API path used continuous random distances, which never produce ties — so the average-rank tie-breaking logic in the Spearman path had no test exercising its non-trivial branch. Added a test using integer-valued distances specifically to force ties, comparing the result against the NumPy path at multiple points.
Both are now covered, and the Array API results match the existing implementation to floating-point precision across sizes and backends.
Verification
| Where | What ran | Result |
|---|---|---|
| NumPy | Full existing PERMANOVA + Mantel test suites | Unchanged, all passing |
| PyTorch (CPU) | New Array API test classes, incl. id-reordering and tie handling | All passing |
| AMD Instinct MI210 (ROCm) | Same Array API test classes, on real GPU hardware | PERMANOVA pseudo-F and Mantel Pearson/Spearman r and p-values match the NumPy result to floating-point precision, including id-reordering happening on the GPU-resident buffer |
The MI210 run is the one that matters most: it's the same code path a user would hit by constructing a DistanceMatrix directly on a GPU and calling permanova or mantel on it, not a simulated stand-in.
What's next
This closes out the first-pass GPU support for PERMANOVA and Mantel — correctness first, matching Igor's stated priority that speed is secondary at this stage. Two things are explicitly deferred as performance follow-ups rather than correctness gaps: each Monte-Carlo permutation currently does a small host-to-device transfer inside the loop, and the Spearman path does a one-time host-side assembly of the ranked matrix. Both are candidates for the next phase of work, alongside the Numba @cuda.jit/@hip.jit kernel backend that the project's roadmap targets for squeezing out real GPU performance rather than just GPU correctness. Distance-matrix class polish (including condensed-form support on GPU) is queued in parallel.
Correctness first, then speed — the same order the CPU phase followed, now repeating on the GPU side.