dgl.graphbolt.expand_indptr

dgl.graphbolt.expand_indptr(indptr, dtype=None, node_ids=None, output_size=None)[source]

Converts a given indptr offset tensor to a COO format tensor. If node_ids is not given, it is assumed to be equal to torch.arange(indptr.size(0) - 1, dtype=dtype, device=indptr.device).

This is equivalent to

if node_ids is None:
    node_ids = torch.arange(len(indptr) - 1, dtype=dtype, device=indptr.device)
return node_ids.to(dtype).repeat_interleave(indptr.diff())
Parameters:

  • indptr (torch.Tensor) – A 1D tensor represents the csc_indptr tensor.

  • dtype (Optional[torch.dtype]) – The dtype of the returned output tensor.

  • node_ids (Optional[torch.Tensor]) – A 1D tensor represents the column node ids that the returned tensor will be populated with.

  • output_size (Optional[int]) – The size of the output tensor. Should be equal to indptr[-1]. Using this argument avoids a stream synchronization to calculate the output shape.

Returns:

The converted COO tensor with values from node_ids.

Return type:

torch.Tensor