GPU indices¶

Present only when the extension was built with GPU support, which is the default. See GPU for the three things that differ from the CPU estimators.

gpu ¶

GPU estimators, present only when the extension was built with them.

ann_search.gpu_available() is the check worth making: it says whether this machine has an adapter, which is the part that varies. The ordinary wheel has GPU support compiled in, so importing this module only fails on a build made with --no-default-features.

Three things differ from the CPU estimators, all of them consequences of the backend rather than choices:

float32 only. WGSL has no float64, so fit casts rather than letting a float64 array fail somewhere inside a kernel. That is a narrowing conversion, and the only one this package performs silently.
No persistence. These indices hold device buffers and sit outside the crate's serialise feature, so save, load and pickle raise. Rebuild.
Manhattan is unavailable on all three.

ExhaustiveGpuIndex ¶

ExhaustiveGpuIndex(
    n_neighbors: int = 15,
    metric: str = "euclidean",
    verbose: bool = False,
)

Bases: _BaseGpuIndex

Brute-force exact search on the device.

Recall is 1 by construction, and on a dataset too large to score on the CPU this is the cheapest way to get ground truth. No build-time knobs: the data goes up, the norms get recorded, and that is the index.

Parameters:

Name	Type	Description	Default
`n_neighbors`	`int`	Neighbours per query, and the default `k` for `kneighbors`.	`15`
`metric`	`str`	`"euclidean"`/`"l2"`, `"sqeuclidean"` or `"cosine"`.	`'euclidean'`
`verbose`	`bool`	Progress to the process stdout, not `sys.stdout`. In Jupyter that lands in the terminal running the kernel.	`False`

IvfGpuIndex ¶

IvfGpuIndex(
    n_neighbors: int = 15,
    metric: str = "euclidean",
    nlist: int | None = None,
    nprobe: int | None = None,
    nquery: int | None = None,
    kmeans_iters: int | None = None,
    kmeans_balanced: bool = False,
    quantise_to_f16: bool = False,
    seed: int = 42,
    verbose: bool = False,
)

Bases: _BaseGpuIndex

Inverted file with k-means and the vectors both on the device.

Trains and queries without a readback, which is what makes it quick, and also what bounds it: the reordered vectors stay resident, so the dataset has to fit in device memory. nquery caps how many queries stage per batch if the upload is what does not fit.

Parameters:

Name	Type	Description	Default
`n_neighbors`	`int`	Neighbours per query, and the default `k` for `kneighbors`.	`15`
`metric`	`str`	`"euclidean"`/`"l2"`, `"sqeuclidean"` or `"cosine"`.	`'euclidean'`
`nlist`	`int \| None`	Number of Voronoi cells to cut the space into. `None` defaults to `sqrt(n)`.	`None`
`nprobe`	`int \| None`	Cells visited per query, the recall knob. `None` defaults to `sqrt(nlist)`. Search-time: override it per call as `index.kneighbors(nprobe=32)`.	`None`
`nquery`	`int \| None`	Queries staged on the device per batch. `None` sizes the batch from `nprobe`, the average cell size and the device's maximum binding, clamped into `100..=20_000`. Lower it if a large `nprobe` overruns the candidate buffer. Search-time.	`None`
`kmeans_iters`	`int \| None`	Lloyd iterations when training the cells. `None` defaults to 50, which is the GPU default and higher than the CPU side's 30: the iterations are cheap once the data is resident.	`None`
`kmeans_balanced`	`bool`	Reseed starved centroids each iteration, RAFT-style. Evens out the posting lists, which matters more here than on the CPU because a straggler cell serialises its whole workgroup.	`False`
`quantise_to_f16`	`bool`	Hold the resident data buffer at fp16. Halves its memory and lifts effective bandwidth on the assignment kernels, at the cost of needing `shader-f16` on the adapter.	`False`
`seed`	`int`	Fixes k-means initialisation.	`42`
`verbose`	`bool`	Progress to the process stdout, not `sys.stdout`. In Jupyter that lands in the terminal running the kernel.	`False`

CagraGpuIndex ¶

CagraGpuIndex(
    n_neighbors: int = 15,
    metric: str = "euclidean",
    graph_degree: int | None = None,
    build_k: int | None = None,
    max_iters: int | None = None,
    n_trees: int | None = None,
    delta: float | None = None,
    rho: float | None = None,
    refine_knn: int | None = None,
    retain_gpu: bool = True,
    beam_width: int | None = None,
    max_beam_iters: int | None = None,
    n_entry_points: int | None = None,
    expand_per_iter: int | None = None,
    seed: int = 42,
    verbose: bool = False,
)

Bases: ExtractKnnMixin, _BaseGpuIndex

CAGRA graph: NN-Descent on the device, pruned, then beam-searched.

The fastest route to a kNN graph here when a GPU is present. beam_width is the recall knob; leaving it None sizes the beam as 2 * max(k, 16), which is usually the right answer.

extract_knn hands back the kNN graph the descent converged on, taken before the CAGRA prune. No kernel runs, and it is capped by graph_degree rather than by the beam.

Unlike every other index in this package, a fitted handle is not safe to query from two threads at once: the beam search memoises its upload of the navigational graph behind a mutable borrow, so concurrent calls serialise. extract_knn is exempt.

Parameters:

Name	Type	Description	Default
`n_neighbors`	`int`	Neighbours per query, and the default `k` for `kneighbors`. Not the graph degree, which is `graph_degree`.	`15`
`metric`	`str`	`"euclidean"`/`"l2"`, `"sqeuclidean"` or `"cosine"`.	`'euclidean'`
`graph_degree`	`int \| None`	Neighbours stored per node in the final graph. `None` defaults to 30. The memory knob, and the ceiling on what `extract_knn` can return.	`None`
`build_k`	`int \| None`	Working degree the descent keeps before pruning. `None` defaults to `1.5 * graph_degree`. Room to manoeuvre, so above `graph_degree` is the point.	`None`
`max_iters`	`int \| None`	Descent iteration cap. `None` defaults to 15.	`None`
`n_trees`	`int \| None`	Random projection trees used to seed the initial graph. `None` defaults to `min(5 + round(n ** 0.25), 20)`.	`None`
`delta`	`float \| None`	Convergence threshold. Descent stops once the fraction of updated edges falls below it. `None` defaults to 0.001.	`None`
`rho`	`float \| None`	Local-join sampling rate. `None` defaults to 1.0, i.e. no sampling. Lower it to trade graph quality for build time.	`None`
`refine_knn`	`int \| None`	Two-hop refinement sweeps after the main loop. `None` defaults to 0, so refinement is off.	`None`
`retain_gpu`	`bool`	Upload the navigational graph at build time rather than on the first query. On by default, because the first query would otherwise pay for it.	`True`
`beam_width`	`int \| None`	Beam width at query time, the recall knob. Search-time.	`None`
`max_beam_iters`	`int \| None`	Iteration cap on the beam search. Search-time.	`None`
`n_entry_points`	`int \| None`	Entry points into the graph per query. Search-time.	`None`
`expand_per_iter`	`int \| None`	Extra neighbours explored per beam iteration, usually 1 to 4. Search-time.	`None`
`seed`	`int`	Fixes the seed graph and the sampling.	`42`
`verbose`	`bool`	Progress to the process stdout, not `sys.stdout`. In Jupyter that lands in the terminal running the kernel.	`False`

Note

The four beam parameters are all-or-nothing. Leave every one of them None and the beam is sized from k: beam_width = 2 * max(k, 16) and max_beam_iters = 3 * beam_width, with 8 entry points and 3 expansions per iteration. Set any one of them and that scaling is off for all four, and the ones still None fall back to flat constants: beam_width = 16, max_beam_iters = 48, n_entry_points = 8, expand_per_iter = 3. So asking for n_entry_points=16 alone on a k=50 query quietly narrows the beam from 100 to 16. If you touch one, set beam_width too.