dgl.graphbolt.index_select

dgl.graphbolt.index_select(tensor, index)

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

返回的张量具有与原始张量（tensor）相同数量的维度。第一维度的大小与索引的长度相同；其他维度的大小与原始张量中的相同。

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

Parameters:

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

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

Returns:

索引后的输入张量，等同于 tensor[index]。

Return type:

torch.Tensor