diag_embed¶

paddle.nn.functional. diag_embed ( input, offset=0, dim1=- 2, dim2=- 1 ) [source]

Creates a tensor whose diagonals of certain 2D planes (specified by dim1 and dim2) are filled by input. By default, a 2D plane formed by the last two dimensions of the returned tensor will be selected.

The argument offset determines which diagonal is generated:

• 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|numpy.ndarray) – The input tensor. Must be at least 1-dimensional. The input data type should be float32, float64, int32, int64.

• offset (int, optional) – Which diagonal to consider. Default: 0 (main diagonal).

• dim1 (int, optional) – The first dimension with respect to which to take diagonal. Default: -2.

• dim2 (int, optional) – The second dimension with respect to which to take diagonal. Default: -1.

Returns

Tensor, the output data type is the same as input data type.

Examples

import paddle
import paddle.nn.functional as F

diag_embed_input = paddle.arange(6)

diag_embed_output1 = F.diag_embed(diag_embed_input)
print(diag_embed_output1)
# Tensor(shape=[6, 6], dtype=int64, place=Place(cpu), stop_gradient=True,
#        [[0, 0, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0, 0],
#         [0, 0, 2, 0, 0, 0],
#         [0, 0, 0, 3, 0, 0],
#         [0, 0, 0, 0, 4, 0],
#         [0, 0, 0, 0, 0, 5]])

diag_embed_output2 = F.diag_embed(diag_embed_input, offset=-1, dim1=0,dim2=1 )
print(diag_embed_output2)
# Tensor(shape=[7, 7], dtype=int64, place=Place(cpu), stop_gradient=True,
#        [[0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0, 0, 0],
#         [0, 0, 2, 0, 0, 0, 0],
#         [0, 0, 0, 3, 0, 0, 0],
#         [0, 0, 0, 0, 4, 0, 0],
#         [0, 0, 0, 0, 0, 5, 0]])

diag_embed_input_2dim = paddle.reshape(diag_embed_input,[2,3])
print(diag_embed_input_2dim)
# Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
#        [[0, 1, 2],
#         [3, 4, 5]])
diag_embed_output3 = F.diag_embed(diag_embed_input_2dim,offset= 0, dim1=0, dim2=2 )
print(diag_embed_output3)
# Tensor(shape=[3, 2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
#        [[[0, 0, 0],
#          [3, 0, 0]],

#         [[0, 1, 0],
#          [0, 4, 0]],

#         [[0, 0, 2],
#          [0, 0, 5]]])