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 ofindex
must be in the range of the dim-th of input. Note thatinput
andindex
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
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)