GPU estimators

Present only in a build with the gpu feature, which is the default. See GPU for what runs on the device and what does not.

UMAPGpu

UMAPGpu(
    n_components: int = 2,
    n_neighbors: int = 15,
    metric: str = "euclidean",
    min_dist: float = 0.5,
    spread: float = 1.0,
    n_epochs: int = 500,
    learning_rate: float = 1.0,
    init: str = "spectral",
    ann: str = "nndescent_gpu",
    optimiser: str = "adam_gpu",
    randomised: bool = False,
    init_range: float | None = None,
    seed: int = 42,
    verbose: int = 0,
    nn_params: NeighbourParamsGpu | None = None,
    graph_params: UmapGraph | None = None,
    optim_params: UmapOptim | None = None,
)

Bases: BaseEmbedding

UMAP with a GPU neighbour search and a GPU Adam optimiser.

"nndescent_gpu" builds a CAGRA graph on the device and is the default for good reason: it is the one that scales. "exhaustive_gpu" gives exact neighbours and is the honest choice for ground truth, at quadratic cost.

Parameters:

Name	Type	Description	Default
n_components	int	Output dimensionality.	2
n_neighbors	int	Neighbours per point.	15
metric	str	"euclidean"/"l2" or "cosine".	'euclidean'
min_dist	float	How tightly points may pack. Fits the repulsion curve with spread.	0.5
spread	float	Scale of the embedding relative to min_dist.	1.0
n_epochs	int	Optimisation epochs.	500
learning_rate	float	Initial learning rate.	1.0
init	str	"spectral", "pca" or "random". Computed on the CPU either way.	'spectral'
ann	str	"nndescent_gpu", "ivf_gpu" or "exhaustive_gpu".	'nndescent_gpu'
optimiser	str	"adam_gpu" keeps the update on the device. The CPU names still work and pull the embedding back each epoch, which is only worth it for debugging.	'adam_gpu'
randomised	bool	Use randomised SVD for the PCA initialisation.	False
init_range	float | None	Scale of the initial coordinates.	None
seed	int	Fixes the initialisation and the negative sampling.	42
verbose	int	0 silent, 1 normal, 2 detailed.	0
nn_params	NeighbourParamsGpu | None	See NeighbourParamsGpu.	None
graph_params	UmapGraph | None	See UmapGraph.	None
optim_params	UmapOptim | None	See UmapOptim.	None

DensMAPGpu

DensMAPGpu(
    n_components: int = 2,
    n_neighbors: int = 15,
    metric: str = "euclidean",
    min_dist: float = 0.5,
    spread: float = 1.0,
    lambda_: float = 2.0,
    n_epochs: int = 500,
    learning_rate: float = 1.0,
    init: str = "spectral",
    ann: str = "nndescent_gpu",
    optimiser: str = "adam_gpu",
    randomised: bool = False,
    init_range: float | None = None,
    seed: int = 42,
    verbose: int = 0,
    nn_params: NeighbourParamsGpu | None = None,
    graph_params: UmapGraph | None = None,
    optim_params: UmapOptim | None = None,
    dens_params: DensParams | None = None,
)

Bases: UMAPGpu

densMAP on the GPU.

With the default "adam_gpu" optimiser the density term runs on the device alongside the rest of the update.

Parameters:

Name	Type	Description	Default
lambda_	float	Weight on the density term. 0 recovers plain UMAP.	2.0
dens_params	DensParams | None	Remaining density knobs. See DensParams.	None

Everything else is as UMAPGpu.

TSNEGpu

TSNEGpu(
    n_components: int = 2,
    perplexity: float = 30.0,
    metric: str = "euclidean",
    n_epochs: int = 1000,
    learning_rate: float | None = None,
    init: str = "pca",
    ann: str = "nndescent_gpu",
    approx: str = "barnes_hut",
    randomised_init: bool = True,
    init_range: float | None = None,
    seed: int = 42,
    verbose: int = 0,
    nn_params: NeighbourParamsGpu | None = None,
    optim_params: TsneOptim | None = None,
)

Bases: BaseEmbedding

t-SNE with a GPU neighbour search.

Only the search moves to the device here; the Barnes-Hut repulsion stays on the CPU. On a dataset where the search dominates that is most of the win, and on one where it does not you should not expect much.

When ann is "nndescent_gpu" and NeighbourParamsGpu.k is left unset, the CAGRA graph degree is backfilled to 3 * perplexity so it is sized for the query t-SNE actually makes.

Parameters:

Name	Type	Description	Default
n_components	int	Output dimensionality. Must be 2.	2
perplexity	float	Effective neighbourhood size.	30.0
metric	str	"euclidean"/"l2" or "cosine".	'euclidean'
n_epochs	int	Optimisation epochs.	1000
learning_rate	float | None	None applies the max(N / 12, 200) heuristic.	None
init	str	"pca", "spectral" or "random".	'pca'
ann	str	"nndescent_gpu", "ivf_gpu" or "exhaustive_gpu".	'nndescent_gpu'
approx	str	Repulsion approximation. See manifolds_rs.TSNE.	'barnes_hut'
randomised_init	bool	Use randomised SVD for the PCA initialisation.	True
init_range	float | None	Scale of the initial coordinates.	None
seed	int	Fixes the initialisation.	42
verbose	int	0 silent, 1 normal, 2 detailed.	0
nn_params	NeighbourParamsGpu | None	See NeighbourParamsGpu.	None
optim_params	TsneOptim | None	See TsneOptim.	None