Module tch::index [−][src]
Expand description
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.