dgl.graphbolt.index_select

dgl.graphbolt.index_select(tensor, index)

Returns a new tensor which indexes the input tensor along dimension dim using the entries in index.

The returned tensor has the same number of dimensions as the original tensor (tensor). The first dimension has the same size as the length of index; other dimensions have the same size as in the original tensor.

When tensor is a pinned tensor and index.is_cuda is True, the operation runs on the CUDA device and the returned tensor will also be on CUDA.

Parameters:

• tensor (torch.Tensor) – The input tensor.

• index (torch.Tensor) – The 1-D tensor containing the indices to index.

Returns:

The indexed input tensor, equivalent to tensor[index].

Return type:

torch.Tensor