torch.linalg.diagonal
function diagonal<S extends Shape, D extends DType, Dev extends DeviceType>(input: Tensor<S, D, Dev>, options?: DiagonalOptions): Tensor<DynamicShape, D, Dev>function diagonal<S extends Shape, D extends DType, Dev extends DeviceType>(input: Tensor<S, D, Dev>, offset: number, dim1: number, dim2: number, options?: DiagonalOptions): Tensor<DynamicShape, D, Dev>Extracts a diagonal or super-diagonal from a tensor.
Returns a new tensor containing the elements along the diagonal of a 2D slice through a tensor. For 2D matrices, extracts main diagonal or super/sub-diagonals. For higher-D tensors, applies to specified 2D planes. Essential for:
- Matrix properties: Extracting diagonal for trace, eigenvalues
- Diagonal matrices: Constructing or extracting diagonal structure
- Band matrices: Extracting bands around diagonal
- Condition number computation: Using diagonal for scaling
- Jacobian analysis: Extracting diagonal of Jacobian or Hessian
- Regularization: Accessing diagonal for matrix modifications
- Main diagonal by default: offset=0 extracts principal diagonal
- Super/sub-diagonals: Positive/negative offset for other diagonals
- Shared storage: Returns view of original data (modifications affect original in some implementations)
- Reduced rank: Output always 1D regardless of input dimensions
- Currently 2D: Full N-D support not yet implemented; use 2D matrices
- 2D limitation: Currently only supports 2D tensors; 3D+ will raise error
- Offset bounds: Very large offset values may return empty tensors
Parameters
inputTensor<S, D, Dev>- Input tensor (must be 2D for current implementation)
optionsDiagonalOptionsoptional- Optional settings for diagonal extraction
Returns
Tensor<DynamicShape, D, Dev>– 1D tensor containing the extracted diagonal elementsExamples
// Extract main diagonal
const A = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]);
torch.linalg.diagonal(A); // [1, 5, 9]
// Extract super-diagonal (above main diagonal)
torch.linalg.diagonal(A, { offset: 1 }); // [2, 6]
// Extract sub-diagonal (below main diagonal)
torch.linalg.diagonal(A, { offset: -1 }); // [4, 8]
// Trace: sum of diagonal
const trace = torch.linalg.diagonal(A).sum();
// Extract diagonal for modification
const A_mat = torch.randn(5, 5);
const diag = torch.linalg.diagonal(A_mat);
const A_modified = A_mat.clone();
// Update diagonal elements (in-place: A_modified[i,i] = new_value)See Also
- PyTorch torch.linalg.diagonal(input, *, offset=0, dim1=-2, dim2=-1)
- ops.diag - Create diagonal matrix from diagonal elements
- trace - Sum of diagonal (uses diagonal internally)