mojolearn.IVFRaBitQIndex#

class mojolearn.IVFRaBitQIndex(n_lists, n_probes, n_neighbors=8, kmeans_n_iters=20, random_state=0)#

Approximate nearest-neighbor index with inverted lists and 1-bit RaBitQ codes (IVF-RaBitQ).

Follows cuVS’s ivf_rabitq (RaBitQ, Gao and Long). fit trains a k-means coarse quantizer, rotates each row’s residual by a randomized Hadamard transform, and keeps its sign bits together with its norm and an inner-product correction factor. search scans the n_probes lists nearest each query and returns the RaBitQ distance estimates, the n_neighbors smallest by (estimate, id). The distances are estimates; pass the candidates to refine for exact distances. Results are reproducible. On a GPU the fitted index stays on the device between searches; each search uploads only its queries.

Parameters:

  • n_lists (int) – Number of inverted lists (k-means clusters). Must be in [1, n_samples].

  • n_probes (int) – Number of lists scanned per query. Required, with no default.

  • n_neighbors (int, default 8) – Number of neighbors search returns per query.

  • kmeans_n_iters (int, default 20) – k-means iterations for the coarse quantizer.

  • random_state (int, default 0) – Seed for the k-means initialization and the rotation’s random signs.

  • numeric_mode (str or None, default None) – Numeric mode for this estimator; None uses the package default set with mojolearn.set_numeric_mode.

Variables:

  • centers_ (ndarray of shape (n_lists, n_features), float32) – Coarse quantizer centers.

  • list_offsets_ (ndarray of shape (n_lists + 1,), int32) – Start of each list in list_indices_.

  • list_indices_ (ndarray of shape (n_samples,), int32) – Original row ids, grouped by list.

  • codes_ (ndarray of shape (n_samples, n_words), int32) – Sign bits of each rotated residual, packed 32 per word; the residual is zero-padded to the next power of two.

  • norms_ (ndarray of shape (n_samples,), float32) – Norm of each residual.

  • ip_factors_ (ndarray of shape (n_samples,), float32) – Inner-product correction factor of each row.

  • n_features_in_ (int) – Number of features seen during fit.

  • n_rows_ (int) – Number of indexed rows.

  • n_lists_ (int) – The fitted value of n_lists.

  • seed_ (int) – The seed used for the rotation.

  • n_candidates_ (ndarray of shape (n_queries,), int32) – Set by search: how many indexed rows each query scanned.

Raises:

  • TypeError – If an integer parameter is not an int.

  • ValueError – If n_lists is out of range.

See also

refine

Re-rank candidates with exact distances.

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> index = mojolearn.IVFRaBitQIndex(n_lists=16, n_probes=4, n_neighbors=20).fit(X)
>>> _, candidates = index.search(X[:10])
>>> distances, indices = mojolearn.refine(X, X[:10], candidates, k=5)

See also

NearestNeighbors, IVFIndex, IVFPQIndex, IVFSQIndex, CagraIndex

User guide: Neighbors and vector search

Choosing an estimator

Methods

fit(X[, y])

Build the index from X.

get_params([deep])

load(path)

Load an index written by save.

numeric_mode_used()

The numeric mode this estimator runs in.

save(path)

Write the fitted index to path as an .npz file.

search(queries[, filter])

Find the approximate nearest indexed rows of each query.

set_params(**params)

Set constructor parameters.

vendor_used()

The GPU backend this estimator runs on.
fit(X, y=None)#

Build the index from X.

Parameters:

  • X (array-like of shape (n_samples, n_features)) – Rows to index, converted to float32.

  • y (None) – Ignored.

Returns:

self (IVFRaBitQIndex) – The fitted index.

search(queries, filter=None)#

Find the approximate nearest indexed rows of each query.

Parameters:

  • queries (array-like of shape (n_queries, n_features)) – Query rows, converted to float32.

  • filter (array-like of bool, shape (n_rows_,), optional) – One flag per indexed row; a row flagged False is never returned (cuVS’s sample filter, applied before scoring).

Returns:

  • distances (ndarray of shape (n_queries, n_neighbors), float32) – Estimated squared Euclidean distances, ascending.

  • indices (ndarray of shape (n_queries, n_neighbors), int32) – Original row ids of the neighbors. A query whose probed lists hold fewer than n_neighbors rows is padded with distance inf and id -1.

Raises:

ValueError – If the index is not fitted, the queries have a different number of features, or filter has the wrong shape or dtype.

classmethod load(path)#

Load an index written by save.

The arrays are read back byte for byte; a file whose arrays have a different dtype, or that was saved by another class, is refused.

Parameters:

path (str or path-like) – File written by save.

Returns:

index – A fitted index of this class, ready to search.

Raises:

ValueError – If the file was saved by another class or does not hold this class’s parameters.

numeric_mode_used()#

The numeric mode this estimator runs in.

Returns:

str – 'identical' or 'fast', read from the compiled binary the estimator uses, not from the argument it was given.

save(path)#

Write the fitted index to path as an .npz file.

The file holds the constructor parameters, the fitted sizes, the index arrays exactly as fit left them, and the numeric mode. Reload it with load; a loaded index searches without rebuilding.

Parameters:

path (str or path-like) – Destination file.

Returns:

path (str or path-like) – The path written.

Raises:

RuntimeError – If the index has not been fitted.

set_params(**params)#

Set constructor parameters. The estimator is rebuilt through its constructor, so every constructor check runs again and any fitted state is dropped, as a new setting invalidates it.

vendor_used()#

The GPU backend this estimator runs on.

Returns:

str – 'metal', 'cuda' or 'hip', read from the compiled binary the estimator uses.