burn::tensor::linalg

Function diag

pub fn diag<B, const D: usize, const DO: usize, K>(
    tensor: Tensor<B, D, K>,
) -> Tensor<B, DO, K>
where
    B: Backend,
    K: TensorKind<B> + BasicOps<B>,
Expand description

Returns the diag of a matrix.

For batched inputs, returns of each matrix in the batch independently.

The diag operation extracts the diagonal elements of the last two dimensions, treating them as the matrix dimensions, while preserving all leading batch dimensions.

§Arguments

  • tensor - The input tensor with at least 2 dimensions.

§Returns

A tensor of rank D - 1, where the last dimension contains the diagonal elements of the input.