rllm.utils.index_to_ptr¶

class rllm.utils.index_to_ptr(index: Tensor, num_nodes: int | None = None)[source]¶

Bases:

Convert a sorted index tensor to a CSR/CSC pointer tensor.

Parameters:

index (Tensor) – The sorted index tensor.
num_nodes (int, optional) – The number of nodes. If None, inferred from index. (default: None)

Example

>>> index = torch.tensor([0, 1, 1, 2, 2, 3])
>>> index_2_ptr(index, 4)
tensor([0, 1, 3, 5, 6])