mojolearn.IVFSQIndex#

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

Approximate nearest-neighbor index with inverted lists and 8-bit scalar quantization (IVF-SQ).

Follows cuVS’s ivf_sq. fit trains the same k-means coarse quantizer as IVFPQIndex, then stores each row’s residual (row minus its list center) quantized to 8 bits per feature, using cuVS’s per-feature range with a 5% margin. search scans the decoded residuals of the n_probes lists nearest each query and returns approximate squared Euclidean 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.
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.
sq_vmin_ (ndarray of shape (n_features,), float32) – Lower end of each feature’s quantization range.
sq_delta_ (ndarray of shape (n_features,), float32) – Quantization step of each feature.
codes_ (ndarray of shape (n_samples, n_features), int32) – 8-bit code of every residual value.
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.
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.

Examples

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

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 (IVFSQIndex) – 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) – Approximate 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.