GPU¶

Four estimators run part of the work on the device: UMAPGpu, DensMAPGpu, TSNEGpu and DREAMSGpu. The backend is wgpu, so there is no CUDA runtime to install and the same wheel works on Metal, Vulkan and DX12.

import manifolds_rs as mf

mf.gpu_available()  # True only if this build has GPU support AND an adapter exists

That one call answers both questions, which is what a caller choosing between UMAP and UMAPGpu actually wants to know.

What actually moves to the device¶

Not everything. The neighbour search and the Adam update run on the GPU; the graph construction and the initialisation stay on the CPU. For t-SNE it depends on approx: with "barnes_hut" the repulsion stays on the CPU, with "fft_3k_gpu" the whole optimisation runs on the device.

mf.TSNEGpu(approx="fft_3k_gpu").fit_transform(X)

"fft_3k_gpu" is FFT-interpolated repulsion with three kernels (q, q^2 dx, q^2 dy) instead of the four-term expansion, with the embedding resident on the device between epochs. "fft" and "fft_3k" run the same FFT schemes on the CPU after a GPU neighbour search.

Otherwise the win depends on whether the neighbour search was your bottleneck. On a few thousand points it will not be, and the transfers will make the GPU path slower than the CPU one. Past a few hundred thousand it usually is. Measure on your own data and hardware before switching; on Apple Silicon in particular the GPU rarely blows the CPU out of the water for these workloads.

float32 only¶

WGSL has no float64. fit casts a float64 array down rather than letting it fail somewhere inside a kernel, and that is the only silent narrowing this package performs anywhere. If that matters for your data, stay on the CPU estimators, which run in whatever precision you hand them.

Reproducibility¶

GPU embeddings are reproducible in structure, not in coordinates. The device neighbour searches are not always bit-stable at scale, and 500 epochs of Adam turn a 0.16% difference in the graph into visibly different coordinates. The clusters, the branches and the separations come back; the exact positions may not. See the guide for the measurements.

The GPU Adam update itself is bit-stable given a fixed graph, so handing in a precomputed graph makes a GPU run reproducible too.

Different neighbour backends¶

The device backends are not the CPU ones with a suffix. "nndescent_gpu" builds a CAGRA graph and is the default because it is the one that scales; "ivf_gpu" and "exhaustive_gpu" are the other two. Passing a CPU-only name like "hnsw" is an error rather than a silent fallback.

Their knobs live in NeighbourParamsGpu, which is a different set from the CPU NeighbourParams: CAGRA build degrees and beam-search budgets have no CPU counterpart.

mf.UMAPGpu(
    n_neighbors=15,
    ann="nndescent_gpu",
    nn_params=mf.NeighbourParamsGpu(k_build=64, beam_width=32),
)

Under TSNEGpu with nndescent_gpu, leaving NeighbourParamsGpu.k unset backfills the graph degree to 3 * perplexity, so it is sized for the query t-SNE actually makes.

Mixing CPU and GPU¶

The neighbour graph is just two arrays, so you can build it wherever it is cheapest and hand it to whichever estimator you like:

ind, dist = mf.knn_graph(X, k=15, ann="hnsw")  # CPU search
embedding = mf.UMAPGpu().fit_transform(X, knn_indices=ind, knn_distances=dist)

That skips the GPU search entirely and leaves only the Adam update on the device, which is occasionally the right split.

Headless machines and CI¶

gpu_available() returns False on a box with no adapter rather than raising, and the GPU estimators are still importable there. That is deliberate: a hosted CI runner can prove the GPU wheel compiles and degrades cleanly without having a device. It cannot prove the kernels are correct. That needs a runner with a GPU.

MLX on Apple Silicon¶

There is a second Apple GPU path, through MLX via mlx-c, and it is not in the published wheel. It links mlx-c dynamically, so a wheel carrying it fails to import on any Mac without the library. Build it yourself:

brew install mlx-c
maturin develop --release --features mlx

mf.mlx_available()  # True only in an mlx build that found mlx-c

MLX is not a separate set of estimators. The CPU estimators gain a few extra names instead:

mf.UMAP(ann="nndescent_mlx", optimiser="adam_mlx").fit_transform(X)
mf.DensMAP(optimiser="adam_mlx").fit_transform(X)
mf.TSNE(ann="nndescent_mlx", approx="fft_3k_mlx").fit_transform(X)
mf.DREAMS(approx="fft_3k_mlx").fit_transform(X)
mf.knn_graph(X, k=15, ann="ivf_mlx")

The neighbour backends are "exhaustive_mlx", "ivf_mlx" and "nndescent_mlx". TSNEGpu and DREAMSGpu take approx="fft_3k_mlx" as well, with the search on wgpu and the optimiser on MLX. den-SNE has no MLX path. Everything MLX runs in float32, whatever you hand it. In a build without MLX these names are refused with a pointer to the build flag.

Is it worth the trouble? On an M1 Max the MLX neighbour search is 1.2x to 1.8x faster than wgpu. UMAP's optimiser ties, and the t-SNE one falls to half of wgpu's speed at a million points, so keep TSNEGpu(approx="fft_3k_gpu") for large t-SNE. See the wgpu vs MLX benchmarks.