Tensor flip along dimension

[rMazeiks]

Thanks for the amazing library! 🦀

Feature description

I would like to flip a tensor along a given dimension. This would be like pytorch's flip function.

Basically it just creates a "mirror" image of the tensor:

Reverse the order of an n-D tensor along given axis in dims.

Parameters

input (Tensor) – the input tensor.

dims (a list or tuple) – axis to flip on

Feature motivation

I want to create a reverse diagonal, but I'm sure there are many more usecases, as this seems like a very basic operation.

Forgive me if there's a way to do this that I missed.

[rMazeiks]

If anyone else runs into this, here's code for creating a reverse-diagonal square matrix:

fn reverse_diagonal<B: Backend>(size: usize) -> Tensor<B, 2> {
    let t = Tensor::zeros([size * size]);
    let indices = Tensor::arange(1..(size + 1)).mul_scalar(size as i64 - 1);
    let values = Tensor::ones([size]);
    t.select_assign(0, indices, values).reshape([size, size])
}

You can left-multiply with this to flip the rows, or right-multiply with it to flip the columns. It should work with any matrix, even non-square ones.

However, this feels like a workaround, and a native flip operation would be nice I think.

[dush-t]

Hi, I'd like to take a shot at this, if nobody is working on it already.

[carrotflakes]

@dush-t
Hi, Are you still working on it?

[antimora]

@carrotflakes it has been a month and we don't we any open PR related. You can give a try. We will review your quickly.

[jungerm2]

It seems there's a bug in this implementation as it relates to transpose:

let device = &Default::default();
type B = burn::backend::wgpu::JitBackend<WgpuRuntime<AutoGraphicsApi, f32, i32>>;
let arr: Tensor<B, 2> = Tensor::from_ints([[1, 2], [3, 4], [5, 6]], device).float();
let arr2 = arr.flip([0]).transpose(); // equivalent to rotating 90 degrees

// This correctly rotates the matrix and will print out
// [[5, 3, 1], [6, 4, 2]] 
println!("{arr2}");

// But the following should yield the same result and doesn't:
let arr2 = arr.transpose().flip([1]); // equivalent to rotating 90 degrees

// This will print out an unexpected result:
// [[5.0, 6.0, 3.0], [4.0, 1.0, 2.0]]
println!("{arr2}");

Can anyone confirm this?