torch.diagflat

torch.diagflat(input, offset=0) → Tensor

• If input is a vector (1-D tensor), then returns a 2-D square tensor with the elements of input as the diagonal.
• If input is a tensor with more than one dimension, then returns a 2-D tensor with diagonal elements equal to a flattened input.

The argument offset controls which diagonal to consider:

• If offset = 0, it is the main diagonal.
• If offset > 0, it is above the main diagonal.
• If offset < 0, it is below the main diagonal.

Parameters

• input (Tensor) – the input tensor.
• offset (int, optional) – the diagonal to consider. Default: 0 (main diagonal).

Examples:

>>> a = torch.randn(3)
>>> a
tensor([-0.2956, -0.9068,  0.1695])
>>> torch.diagflat(a)
tensor([[-0.2956,  0.0000,  0.0000],
        [ 0.0000, -0.9068,  0.0000],
        [ 0.0000,  0.0000,  0.1695]])
>>> torch.diagflat(a, 1)
tensor([[ 0.0000, -0.2956,  0.0000,  0.0000],
        [ 0.0000,  0.0000, -0.9068,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  0.1695],
        [ 0.0000,  0.0000,  0.0000,  0.0000]])

>>> a = torch.randn(2, 2)
>>> a
tensor([[ 0.2094, -0.3018],
        [-0.1516,  1.9342]])
>>> torch.diagflat(a)
tensor([[ 0.2094,  0.0000,  0.0000,  0.0000],
        [ 0.0000, -0.3018,  0.0000,  0.0000],
        [ 0.0000,  0.0000, -0.1516,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  1.9342]])