dgl.graphbolt.expand_indptr

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

将给定的indptr偏移张量转换为COO格式张量。如果未提供node_ids,则假定其等于torch.arange(indptr.size(0) - 1, dtype=dtype, device=indptr.device)。

这相当于

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) – 一个一维张量,表示csc_indptr张量。

  • dtype (可选[torch.dtype]) – 返回输出张量的数据类型。

  • node_ids (可选[torch.Tensor]) – 一个一维张量,表示返回的张量将被填充的列节点ID。

  • output_size (可选[int]) – 输出张量的大小。应等于 indptr[-1]。使用此参数可以避免流同步以计算输出形状。

Returns:

转换后的COO张量,其值来自node_ids。

Return type:

torch.Tensor