Making scikit-bio's Distance Matrices GPU-Capable

Where things stand

The first phase of the project — replacing selected Cython hot paths with Numba CPU implementations — has landed upstream. PERMANOVA and Mantel now have an optional Numba compute engine, selectable per call, with numerical agreement asserted against the existing implementations and coverage measured with the JIT disabled.

With the CPU story in place, the project moved to the part that motivated it in the first place: GPU support. The goal is not just "run faster on a GPU," but something more structural — let scientific data that is already on the GPU stay there as it passes through scikit-bio's distance functions.

Dispatch by where the data lives

The guiding idea is the Python Array Interface Protocol: the execution path is chosen by where the input array lives, not by a manual flag.

Downgrading from GPU to CPU is allowed when no suitable GPU backend is present; silently promoting CPU data onto the GPU is not. The Array API (xp) layer is the safety net that runs correctly on any compliant backend — NumPy, CuPy, PyTorch, or JAX — from a single implementation.

A GPU-capable DistanceMatrix

For a GPU-resident distance matrix to reach the functions, the DistanceMatrix class itself had to be able to hold a non-NumPy buffer. Previously it coerced everything to NumPy at construction, which quietly forced a copy back to the host.

This first-class class support is the foundation the function-level GPU work builds on, and it went up as a focused pull request against the project's development branch.

PERMANOVA & Mantel on the Array API

Both statistics were given an Array API compute path so the whole calculation — the test statistic and the Monte-Carlo permutation distribution — runs in the array namespace, keeping the matrix on its device.

A couple of details made this exact rather than approximate:

• PERMANOVA's within-group sum of squares and pseudo-F reduce to a handful of Array API reductions, reusing the same permutation scheme as the reference so p-values match exactly for a given seed.
• Mantel's normalisation constants are permutation-invariant, so each permuted correlation reuses them; the Spearman variant needed an on-device average-tie ranking that reproduces scipy's ranking element-for-element, including ties.

Across sizes and backends the Array API results match the existing implementation to floating-point precision, with identical p-values.

One source, two vendors

The most satisfying result: the same source runs correctly on both NVIDIA and AMD GPUs. The Array API compute is vendor-neutral, and the fused GPU kernels use one code body that a single-source toolchain compiles for each vendor's hardware.

Engineering notes

• Portability bites in small places. An operation that worked on NumPy needed an explicit axis on another backend — the kind of gap you only find by actually running on multiple backends, not by reading docs.
• Device placement is invisible until it isn't. Intermediate index arrays have to be created on the same device as the data; on a single-device backend this is free, but on a CPU/GPU split it has to be explicit. Testing only on CPU hides this — verifying on a real GPU surfaced it.
• A fallback path is worth keeping. Making the Array API layer the safety net means correctness never depends on a specialised kernel being available, which keeps the fast paths honest.
• Match the codebase, don't reinvent it. Pushing the backend-awareness down into shared helper functions kept the class methods clean and reused existing structure instead of duplicating logic inline.

What's next

This DistanceMatrix support has since merged, and PERMANOVA and Mantel now build on it with their own GPU-resident Array API path — covered in the next update. From here the near-term arc continues toward improving actual GPU performance (a Numba @cuda.jit/@hip.jit kernel backend) and bringing the ordination path (PCoA and its centering step) along too. Work proceeds as stacked pull requests against the development branch so review and the next step can happen in parallel.