Attention
The vector search and clustering algorithms in RAFT are being migrated to a new library dedicated to vector search called cuVS. We will continue to support the vector search algorithms in RAFT during this move, but will no longer update them after the RAPIDS 24.06 (June) release. We plan to complete the migration by RAPIDS 24.08 (August) release.
Matrix#
This page provides pylibraft
class references for the publicly-exposed elements of the pylibraft.matrix
package.
- pylibraft.matrix.select_k(dataset, k=None, distances=None, indices=None, select_min=True, handle=None)[source]#
Selects the top k items from each row in a matrix
- Parameters:
- datasetarray interface compliant matrix, row-major layout,
shape (n_rows, dim). Supported dtype [float]
- kint
Number of items to return for each row. Optional if indices or distances arrays are given (in which case their second dimension is k).
- distancesOptional array interface compliant matrix shape
(n_rows, k), dtype float. If supplied, distances will be written here in-place. (default None)
- indicesOptional array interface compliant matrix shape
(n_rows, k), dtype int64_t. If supplied, neighbor indices will be written here in-place. (default None)
- select_min:bool
Whether to select the minimum or maximum K items
- handleOptional RAFT resource handle for reusing CUDA resources.
If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling
handle.sync()
before accessing the output.
- Returns:
- distances: array interface compliant object containing resulting distances
shape (n_rows, k)
- indices: array interface compliant object containing resulting indices
shape (n_rows, k)
Examples
>>> import cupy as cp
>>> from pylibraft.matrix import select_k
>>> n_features = 50 >>> n_rows = 1000
>>> queries = cp.random.random_sample((n_rows, n_features), ... dtype=cp.float32) >>> k = 40 >>> distances, ids = select_k(queries, k) >>> distances = cp.asarray(distances) >>> ids = cp.asarray(ids)