torch.Tensor.Tensor.index_select
Tensor.index_select<Dim extends number, IndexLen extends number>(dim: Dim, index: Tensor<readonly [IndexLen]>): Tensor<DynamicShape, D, Dev>Tensor.index_select(dim: number, index: AnyTensor): Tensor<DynamicShape, D, Dev>Selects elements from a tensor along a dimension using indices.
Returns a new tensor with elements selected from this tensor along the specified dimension, using the given 1D index tensor. Useful for:
- Gathering elements by indices
- Sampling subsets of data
- Advanced indexing operations
- Reordering elements along dimensions
- Selecting subset of classes/features
- Indices can have duplicates (elements can be selected multiple times)
- Indices must be in valid range [0, size_along_dim)
- Can use negative dimension (e.g., -1 for last dimension)
- Similar to gather() but only along one dimension
Parameters
dimDim- Dimension along which to select
indexTensor<readonly [IndexLen]>- 1D tensor of indices to select (must be 1D)
Returns
Tensor<DynamicShape, D, Dev>– Tensor with selected elements, shape = input.shape except dim = index.shape[0]Examples
const x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]);
const idx = torch.tensor([0, 2]);
// Select rows 0 and 2
x.index_select(0, idx);
// [[1, 2, 3],
// [7, 8, 9]]
// Select columns 0 and 2
x.index_select(1, idx);
// [[1, 3],
// [4, 6],
// [7, 9]]
// Select multiple copies of same indices
const indices = torch.tensor([1, 0, 1, 0]); // Select 4 elements
x.index_select(0, indices); // Select rows [1, 0, 1, 0]
// [[4, 5, 6],
// [1, 2, 3],
// [4, 5, 6],
// [1, 2, 3]]
// Select features from batch
const features = torch.randn(32, 256); // 32 samples, 256 features
const important_feat = torch.tensor([0, 5, 10, 15]);
const selected = features.index_select(1, important_feat); // [32, 4]See Also
- PyTorch torch.index_select()
- gather - More general element selection with multi-dim indices
- take - Select elements from flattened tensor
- select - Select single element or slice