Skip to main content

movedim

Function movedim 

Source
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 tensor
  • src - Source axis position
  • dst - 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]));