On this page
torch.diag
torch.diag(input, diagonal=0, *, out=None) → Tensor
-
- If
input
is a vector (1-D tensor), then returns a 2-D square tensor with the elements ofinput
as the diagonal. - If
input
is a matrix (2-D tensor), then returns a 1-D tensor with the diagonal elements ofinput
.
The argument
diagonal
controls which diagonal to consider:- If
diagonal
= 0, it is the main diagonal. - If
diagonal
> 0, it is above the main diagonal. - If
diagonal
< 0, it is below the main diagonal.
- Parameters
- Keyword Arguments
-
out (Tensor, optional) – the output tensor.
See also
torch.diagonal()
always returns the diagonal of its input.torch.diagflat()
always constructs a tensor with diagonal elements specified by the input.Examples:
Get the square matrix where the input vector is the diagonal:
>>> a = torch.randn(3) >>> a tensor([ 0.5950,-0.0872, 2.3298]) >>> torch.diag(a) tensor([[ 0.5950, 0.0000, 0.0000], [ 0.0000,-0.0872, 0.0000], [ 0.0000, 0.0000, 2.3298]]) >>> torch.diag(a, 1) tensor([[ 0.0000, 0.5950, 0.0000, 0.0000], [ 0.0000, 0.0000,-0.0872, 0.0000], [ 0.0000, 0.0000, 0.0000, 2.3298], [ 0.0000, 0.0000, 0.0000, 0.0000]])
Get the k-th diagonal of a given matrix:
>>> a = torch.randn(3, 3) >>> a tensor([[-0.4264, 0.0255,-0.1064], [ 0.8795,-0.2429, 0.1374], [ 0.1029,-0.6482,-1.6300]]) >>> torch.diag(a, 0) tensor([-0.4264,-0.2429,-1.6300]) >>> torch.diag(a, 1) tensor([ 0.0255, 0.1374])
- If
© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/generated/torch.diag.html