pylibraft-cu12

import cupy as cp from raft.neighbors import NearestNeighbors import sys # Ensure CuPy is installed for GPU array handling and a CUDA device is available try: if cp.cuda.runtime.getDeviceCount() == 0: print("Error: No CUDA device found. RAFT requires a CUDA-enabled GPU.") sys.exit(1) except cp.cuda.runtime.CUDARuntimeError as e: print(f"Error initializing CuPy or CUDA runtime: {e}") print("Please ensure CuPy is installed correctly (e.g., `cupy-cuda12x`) and your CUDA environment is set up.") sys.exit(1) # Generate some random GPU data using CuPy n_samples = 100 n_features = 10 X = cp.random.rand(n_samples, n_features, dtype=cp.float32) # Create a NearestNeighbors model nn = NearestNeighbors(n_neighbors=5, metric='euclidean', output_type='cupy') # Fit the model to the data nn.fit(X) # Query for nearest neighbors to the same data distances, indices = nn.kneighbors(X) print("Distances shape:", distances.shape) print("Indices shape:", indices.shape) print("\nFirst 5 distances (sample 0):\n", distances[0, :5]) print("\nFirst 5 indices (sample 0):\n", indices[0, :5])