index_select

→Tensor

index_select(input: Tensor, dim: int, index: Tensor)

source

Select slices of input along dim given index.

Unlike gather, the index is a 1-D tensor and the output preserves all axes of input — only dim changes size to len(index).

Parameters

inputTensor

Source tensor.

dimint

Axis along which to select.

indexTensor

1-D int64 tensor whose entries lie in [0, input.size(dim)).

Returns

Tensor

Tensor with the same shape as input except size len(index) along dim.

Notes

Equivalent to input[..., index, ...] with the slice placed at position dim. Reordering, duplication and subsetting are all supported in one call.

Examples

>>> import lucid
>>> x = lucid.tensor([[1, 2], [3, 4], [5, 6]])
>>> lucid.index_select(x, dim=0, index=lucid.tensor([0, 2]))
Tensor([[1, 2],
        [5, 6]])