dgl.graphbolt.expand_indptr

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

将给定的 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())
参数:
  • indptr (torch.Tensor) – 表示 csc_indptr 张量的 1D 张量。

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

  • node_ids (Optional[torch.Tensor]) – 表示将用于填充返回张量的列节点 id 的 1D 张量。

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

返回:

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

返回类型:

torch.Tensor