oneflow.index_select

oneflow.index_select(dim, index)Tensor

Select values along an axis specified by dim.

index must be an Int32 Tensor with 1-D. dim must be in the range of input Dimensions. value of index must be in the range of the dim-th of input. Note that input and index do not broadcast against each other.

The interface is consistent with PyTorch. The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.index_select.html.

Parameters
  • input (Tensor) – the source tensor

  • dim (int) – the axis along which to index

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

For example:

>>> import oneflow as flow
>>> input = flow.tensor([[1,2,3],[4,5,6]], dtype=flow.int32)
>>> input
tensor([[1, 2, 3],
        [4, 5, 6]], dtype=oneflow.int32)
>>> index = flow.tensor([0,1], dtype=flow.int64)
>>> output = flow.index_select(input, 1, index)
>>> output
tensor([[1, 2],
        [4, 5]], dtype=oneflow.int32)
>>> output = input.index_select(1, index)
>>> output
tensor([[1, 2],
        [4, 5]], dtype=oneflow.int32)