dgl.graphbolt.index_select

dgl.graphbolt.index_select(tensor, index)[source]

返回一个新的张量，该张量使用 index 中的条目沿 dim 维度索引输入张量。

返回的张量与原始张量 (tensor) 具有相同的维度数。第一个维度的尺寸与 index 的长度相同；其他维度的尺寸与原始张量相同。

当 tensor 是一个 pinned 张量且 index.is_cuda 为 True 时，操作将在 CUDA 设备上运行，返回的张量也将位于 CUDA 上。

参数:

  • tensor (torch.Tensor) – 输入张量。

  • index (torch.Tensor) – 包含用于索引的索引的一维张量。

返回值:

索引后的输入张量，等同于 tensor[index]。如果 index 位于 pinned memory 中，则结果也将被放入 pinned memory。

返回类型:

torch.Tensor