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.


This page provides pylibraft class references for the publicly-exposed elements of the pylibraft.random package.

pylibraft.random.rmat(out, theta, r_scale, c_scale, seed=12345, handle=None)[source]#

Generate RMAT adjacency list based on the input distribution.

out: CUDA array interface compliant matrix shape (n_edges, 2). This will

contain the src/dst node ids stored consecutively like a pair.

theta: CUDA array interface compliant matrix shape

(max(r_scale, c_scale) * 4) This stores the probability distribution at each RMAT level

r_scale: log2 of number of source nodes
c_scale: log2 of number of destination nodes
seed: random seed used for reproducibility
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.


>>> import cupy as cp
>>> from pylibraft.common import Handle
>>> from pylibraft.random import rmat
>>> n_edges = 5000
>>> r_scale = 16
>>> c_scale = 14
>>> theta_len = max(r_scale, c_scale) * 4
>>> out = cp.empty((n_edges, 2), dtype=cp.int32)
>>> theta = cp.random.random_sample(theta_len, dtype=cp.float32)
>>> # A single RAFT handle can optionally be reused across
>>> # pylibraft functions.
>>> handle = Handle()
>>> rmat(out, theta, r_scale, c_scale, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()