diag_embed

paddle. 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

>>> diag_embed_input = paddle.arange(6)

>>> diag_embed_output1 = paddle.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 = paddle.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 = paddle.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]]])