Newbie question: what is the corresponding method of mat[:, 1, :] with mat as a 3-D matrix in pyTorch?

[Magicloud]

I looked over https://docs.rs/burn/latest/burn/prelude/struct.Tensor.html, and did not find a way to directly do this.

[nathanielsimard]

I think it's a simple slice or slice_assign method call that you need. When using this function you need to specify the range included in every dimension:

let [d1, d2, d3] = mat.shape();
// to read the matrix, like `x = mat[:, 1, :]`
let x = mat.slice(0..d1, 0..1, 0..d3);

// to write to the matrix, like `mat[:, 1, :] = value` but it returns the results instead of mutating the tensor inplace.
let x = mat.slice_assign(0..d1, 0..1, 0..d3, value);

[Magicloud]

Thanks. I think my confusion is because of Python's type.

The result of mat[:, 1, :] being used later appeared to be a 2-D matrix. And slice' type indicates that a matrix with the same dimension is returned.

So I added a squeeze after the slicing.