MLX indices¶

Present only in a source build with the mlx feature. See GPU for how to build it and what differs.

mlx ¶

MLX estimators, present only in a source build with the mlx feature.

Experimental and Apple Silicon only. MLX drives Metal directly through mlx-c rather than going through wgpu, and comes out faster than the ann_search.gpu counterparts at the same recall; the GPU page of the docs has the numbers. The published wheel does not carry them, since it would have to bundle mlx-c. Build from source instead::

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

ann_search.mlx_available() says whether the build picked it up.

Same constraints as the GPU estimators: float32 only (fit casts), no persistence, no Manhattan. One more on top: MLX streams are tied to the thread that created them, so a fitted index raises if used from any other Python thread. Host-side work inside the MLX paths uses rayon's global pool and does not follow ann_search.set_num_threads.

ExhaustiveMlxIndex ¶

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

Bases: _BaseMlxIndex

Brute-force exact search on MLX: its GEMM plus a Metal top-k kernel.

Recall is 1 by construction.

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`	Device versus host time split, to the process stdout.	`False`

IvfMlxIndex ¶

IvfMlxIndex(
    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,
    seed: int = 42,
    verbose: bool = False,
)

Bases: _BaseMlxIndex

Inverted file on MLX: k-means, the vectors and the cell scan on device.

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. `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 per device tile. `None` lets the library size it. Search-time.	`None`
`kmeans_iters`	`int \| None`	Lloyd iterations. `None` defaults to 30. Every one runs; MLX k-means has no early stop.	`None`
`kmeans_balanced`	`bool`	Reseed starved centroids each iteration.	`False`
`seed`	`int`	Fixes k-means initialisation.	`42`
`verbose`	`bool`	Progress to the process stdout, not `sys.stdout`.	`False`

CagraMlxIndex ¶

CagraMlxIndex(
    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,
    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, _BaseMlxIndex

CAGRA graph on MLX: NN-Descent, the prune and the beam search on device.

Same parameters and beam semantics as CagraGpuIndex, minus retain_gpu: the MLX index is resident from the build on. External queries pick their entry points through a random projection forest on the device.

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. `None` defaults to 30.	`None`
`build_k`	`int \| None`	Working degree before pruning. `None` defaults to `1.5 * graph_degree`.	`None`
`max_iters`	`int \| None`	Descent iteration cap. `None` defaults to 15.	`None`
`n_trees`	`int \| None`	Random projection trees for the initial graph. `None` defaults to `min(5 + round(n ** 0.25), 20)`.	`None`
`delta`	`float \| None`	Convergence threshold. `None` defaults to 0.001.	`None`
`rho`	`float \| None`	Local-join sampling rate. `None` defaults to 1.0.	`None`
`refine_knn`	`int \| None`	Two-hop refinement sweeps. `None` defaults to 0.	`None`
`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 per query. Search-time.	`None`
`expand_per_iter`	`int \| None`	Beam entries expanded per iteration. Search-time.	`None`
`seed`	`int`	Fixes the seed graph and the sampling.	`42`
`verbose`	`bool`	Progress to the process stdout, not `sys.stdout`.	`False`

Note

The four beam parameters are all-or-nothing, exactly as on CagraGpuIndex. Leave them all None and the beam is sized from k (beam_width = 2 * max(k, 16), max_beam_iters = 3 * beam_width). Set any one and the rest fall back to flat constants: 16, 48, 8 entry points and 3 expansions. If you touch one, set beam_width too.