pub fn movedim<T>(
arr: &TypedTensor<T>,
src: usize,
dst: usize,
) -> TypedTensor<T>where
T: TensorScalar + Copy,Expand description
Move axis from position src to position dst
This is equivalent to numpy.moveaxis or libsparseir’s movedim. The other axes keep their order.
§Arguments
arr- Input tensorsrc- Source axis positiondst- Destination axis position
§Returns
A new tensor with the axes permuted
§Panics
Panics if src or dst is not an axis of arr.
§Example
use sparse_ir::TypedTensor;
use sparse_ir::sampling::movedim;
// A 4D tensor with shape (2, 3, 4, 5) and entries 1000 i + 100 j + 10 k + l
let shape = [2usize, 3, 4, 5];
let data: Vec<f64> = (0..120)
.map(|lin| {
let (i, j, k, l) = (lin % 2, lin / 2 % 3, lin / 6 % 4, lin / 24);
(1000 * i + 100 * j + 10 * k + l) as f64
})
.collect();
let arr = TypedTensor::from_vec_col_major(shape.to_vec(), data).unwrap();
// movedim(arr, 0, 2) moves axis 0 to position 2
let moved = movedim(&arr, 0, 2);
// Result shape: (3, 4, 2, 5) with axes permuted as [1, 2, 0, 3]
assert_eq!(moved.shape(), &[3, 4, 2, 5]);
// Element [2, 3, 1, 4] of the result is element [1, 2, 3, 4] of arr
let at = |t: &TypedTensor<f64>, idx: [usize; 4]| {
let s = t.shape();
t.host_data().unwrap()[idx[0] + s[0] * (idx[1] + s[1] * (idx[2] + s[2] * idx[3]))]
};
assert_eq!(at(&moved, [2, 3, 1, 4]), at(&arr, [1, 2, 3, 4]));