libraft-cu12

# libraft-cu12 provides the C++ core for RAFT. # Python users typically interact with RAFT functionality via the `pylibraft` package. # Ensure pylibraft-cu12 is installed (it depends on libraft-cu12): # pip install pylibraft-cu12 --extra-index-url https://pypi.nvidia.com import cupy as cp from pylibraft.neighbors import ivf_pq # Example using pylibraft, which relies on libraft-cu12 n_samples = 5000 n_features = 50 # Generate some random data on GPU using CuPy dataset = cp.random.random_sample((n_samples, n_features), dtype=cp.float32) queries = cp.random.random_sample((100, n_features), dtype=cp.float32) # Build an IVF-PQ index index_params = ivf_pq.IndexParams( n_lists=1024, metric="sqeuclidean", pq_dim=10 ) index = ivf_pq.build(index_params, dataset) # Search the index k = 5 search_params = ivf_pq.SearchParams(n_probes=20) distances, neighbors = ivf_pq.search(search_params, index, queries, k) print(f"Distances shape: {distances.shape}") print(f"Neighbors shape: {neighbors.shape}")