megengine.functional.diag#

diag(inp, k=0)[source]#

Extract a diagonal or construct a diagonal tensor.

If inp is a 1D tensor, then returns a 2D tensor with the elements of inp as the diagonal. If inp is a 2D tensor, then returns a 1D tensor with the diagonal elements of inp.

Parameters:

inp – input tensor.
k (int) – diagonal in consider. Use \(k=0\) for the main diagonal, \(k>0\) for diagonals above the main diagonal, and \(k<0\) for diagonals below the main diagonal.

See also

If you want to create a identity matrix, see eye.

Return type:: Tensor
Returns:: the extracted diagonal or constructed diagonal tensor.

Examples

Input is a 1D tensor:

>>> F.diag(Tensor([1, 2, 3]))
Tensor([[1 0 0]
 [0 2 0]
 [0 0 3]], dtype=int32, device=xpux:0)
>>> F.diag(Tensor([1, 2, 3]), k=1)
Tensor([[0 1 0 0]
 [0 0 2 0]
 [0 0 0 3]
 [0 0 0 0]], dtype=int32, device=xpux:0)

Input is a 2D tensor:

>>> x = F.arange(9).reshape(3, 3)
>>> x
Tensor([[0. 1. 2.]
 [3. 4. 5.]
 [6. 7. 8.]], device=xpux:0)
>>> F.diag(x)
Tensor([0. 4. 8.], device=xpux:0)

Get the k-th diagonal of a given matrix:

>>> F.diag(x, k=1)
Tensor([1. 5.], device=xpux:0)
>>> F.diag(x, k=-1)
Tensor([3. 7.], device=xpux:0)