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 siz
(tensor, index)
| 213 | |
| 214 | |
| 215 | def index_select(tensor, index): |
| 216 | """Returns a new tensor which indexes the input tensor along dimension dim |
| 217 | using the entries in index. |
| 218 | |
| 219 | The returned tensor has the same number of dimensions as the original tensor |
| 220 | (tensor). The first dimension has the same size as the length of index; |
| 221 | other dimensions have the same size as in the original tensor. |
| 222 | |
| 223 | When tensor is a pinned tensor and index.is_cuda is True, the operation runs |
| 224 | on the CUDA device and the returned tensor will also be on CUDA. |
| 225 | |
| 226 | Parameters |
| 227 | ---------- |
| 228 | tensor : torch.Tensor |
| 229 | The input tensor. |
| 230 | index : torch.Tensor |
| 231 | The 1-D tensor containing the indices to index. |
| 232 | |
| 233 | Returns |
| 234 | ------- |
| 235 | torch.Tensor |
| 236 | The indexed input tensor, equivalent to tensor[index]. If index is in |
| 237 | pinned memory, then the result is placed into pinned memory as well. |
| 238 | """ |
| 239 | assert index.dim() == 1, "Index should be 1D tensor." |
| 240 | return torch.ops.graphbolt.index_select(tensor, index) |
| 241 | |
| 242 | |
| 243 | def etype_tuple_to_str(c_etype): |
no test coverage detected