Implementation:LaurentMazare Tch rs Tensor Indexing

Knowledge Sources	LaurentMazare_Tch_rs
Domains	Tensor_Operations, Indexing
Last Updated	2026-02-08 00:00 GMT

Overview

Concrete tool for indexing into tensors using Rust-native syntax provided by the tch-rs library.

Description

The Tensor_Indexing module implements the IndexOp trait, which provides the .i() method for tensor indexing. It defines the TensorIndexer enum with four variants (Select, Narrow, IndexSelect, InsertNewAxis) and the NewAxis marker struct. The module supports integer indexing, range-based slicing (all standard Rust range types), tensor-based index selection, and new axis insertion. The IndexOp trait is implemented for single indices and tuples of up to 7 index elements, enabling multi-dimensional indexing in a single call.

Usage

Import the IndexOp trait to access the .i() method on any Tensor. The method accepts integers, ranges, NewAxis, tensor references, and tuples combining these types.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: src/tensor/index.rs
• Lines: 1-372

Signature

pub struct NewAxis;

pub enum TensorIndexer {
    Select(i64),
    Narrow(Bound<i64>, Bound<i64>),
    IndexSelect(Tensor),
    InsertNewAxis,
}

pub trait IndexOp<T> {
    fn i(&self, index: T) -> Tensor;
    fn f_i(&self, index: T) -> Result<Tensor>;
}

The IndexOp trait is implemented for Tensor with the following index types:

• Single index: any A: Into<TensorIndexer>
• 1-tuple through 7-tuple of types implementing Into<TensorIndexer>

Each tuple implementation delegates to the internal f_indexer method:

impl Tensor {
    fn f_indexer(&self, index_spec: &[TensorIndexer]) -> Result<Tensor> { ... }
}

Import

use tch::{IndexOp, NewAxis, Tensor};

I/O Contract

Inputs

Name	Type	Required	Description
self	&Tensor	Yes	The tensor to index into
index	T: Into<TensorIndexer> (or tuple thereof)	Yes	Index specification: integer (i64), range (e.g. 1..3, .., 1..), NewAxis, &Tensor, Vec<i64>, or &[i64]

Outputs

Name	Type	Description
result	Tensor	The indexed sub-tensor (shares underlying storage with the original)

Errors

Condition	Error Type	Description
Too many indices	TchError::Shape	Number of non-NewAxis indices exceeds tensor dimensions
Invalid index tensor kind	TchError::Kind	IndexSelect tensor must be Int64, Int16, Int8, or Int
Multi-dimensional index tensor	TchError::Shape	IndexSelect tensors must be 1-dimensional

Supported Index Types

The following types implement From<T> for TensorIndexer:

Type	Variant	Description
i64	Select(i64)	Selects a single element along a dimension
Range<i64>	Narrow(Bound, Bound)	Narrows a dimension (e.g. 1..3)
RangeFrom<i64>	Narrow(Bound, Bound)	Narrows from start (e.g. 1..)
RangeFull	Narrow(Bound, Bound)	Selects all elements (i.e. ..)
RangeInclusive<i64>	Narrow(Bound, Bound)	Narrows inclusive (e.g. 1..=3)
RangeTo<i64>	Narrow(Bound, Bound)	Narrows up to end (e.g. ..3)
RangeToInclusive<i64>	Narrow(Bound, Bound)	Narrows up to and including end (e.g. ..=3)
&[i64]	IndexSelect(Tensor)	Selects specific indices from a dimension
Vec<i64>	IndexSelect(Tensor)	Selects specific indices from a dimension
&Tensor	IndexSelect(Tensor)	Selects indices specified by a tensor
NewAxis	InsertNewAxis	Inserts a new dimension of size 1

Usage Examples

use tch::{IndexOp, NewAxis, Tensor};

let tensor = Tensor::from_slice(&[1, 2, 3, 4, 5, 6]).view([2, 3]);

// Integer indexing: select row 1
let t = tensor.i(1);
// Result: [4, 5, 6]

// Full range on first dim, specific index on second
let t = tensor.i((.., -2));
// Result: [2, 5]

// Range-based slicing
let t = tensor.i((.., 1..));
// Result shape: [2, 2], values: [2, 3, 5, 6]

let t = tensor.i((..1, ..));
// Result shape: [1, 3], values: [1, 2, 3]

let t = tensor.i((.., 1..=2));
// Result shape: [2, 2], values: [2, 3, 5, 6]

// NewAxis insertion
let t = tensor.i((NewAxis,));
// Result shape: [1, 2, 3]

let t = tensor.i((.., .., NewAxis));
// Result shape: [2, 3, 1]

// Fallible version
let t = tensor.f_i(0)?;  // Returns Result<Tensor>

Related Pages

• Principle:LaurentMazare_Tch_rs_Tensor_Indexing_Operations