Module tch::index

Indexing operations

This module defines the i indexing operation. This can be used in various scenarios.

Using an integer index returns the slice obtained by selecting elements with the specified index. Negative values can be used for the index, and .. can be used to get all the indexes from a given dimension.

use crate::tch::{IndexOp, Tensor};
let tensor = Tensor::of_slice(&[1, 2, 3, 4, 5, 6]).view((2, 3));
let t = tensor.i(1);
let t = tensor.i((.., -2));

Indexes like 1.., ..1, or 1..2, can be used to narrow a dimension.

use crate::tch::{IndexOp, Tensor};
let tensor = Tensor::of_slice(&[1, 2, 3, 4, 5, 6]).view((2, 3));
let t = tensor.i((.., 1..));
assert_eq!(t.size(), [2, 2]);
assert_eq!(Vec::<i64>::from(t.contiguous().view(-1)), [2, 3, 5, 6]);
let t = tensor.i((..1, ..));
assert_eq!(t.size(), [1, 3]);
assert_eq!(Vec::<i64>::from(t.contiguous().view(-1)), [1, 2, 3]);
let t = tensor.i((.., 1..2));
assert_eq!(t.size(), [2, 1]);
assert_eq!(Vec::<i64>::from(t.contiguous().view(-1)), [2, 5]);
let t = tensor.i((.., 1..=2));
assert_eq!(t.size(), [2, 2]);
assert_eq!(Vec::<i64>::from(t.contiguous().view(-1)), [2, 3, 5, 6]);

The NewAxis index can be used to insert a dimension.

use crate::tch::{IndexOp, NewAxis, Tensor};
let tensor = Tensor::of_slice(&[1, 2, 3, 4, 5, 6]).view((2, 3));
let t = tensor.i((NewAxis,));
assert_eq!(t.size(), &[1, 2, 3]);
let t = tensor.i((.., .., NewAxis));
assert_eq!(t.size(), &[2, 3, 1]);

Unlike NumPy, the i operation does not support advanced indexing. The result can be different from NumPy with same set of arguments. For example, tensor.i(..1, vec![0, 3], vec![2, 1, 3]) does narrowing on first dimension, and index selection on second and third dimensions. The analogous NumPy indexing array[:1, [0, 3], [2, 1, 3]] throws shape mismatch error due to advanced indexing rule. Another distinction is that i guarantees the input and result tensor shares the same underlying storage, while NumPy may copy the tensor in certain scenarios.

Structs

NewAxis

Enums

TensorIndexer

Traits

IndexOp