paddle.fluid.layers.trace(input, offset=0, dim1=0, dim2=1, out=None, name=None)[source]

This OP computes the sum along diagonals of the input tensor.

If input is 2D, returns the sum of diagonal.

If input has larger dimensions, then returns an tensor of diagonals sum, diagonals be taken from the 2D planes specified by dim1 and dim2. By default, the 2D planes formed by the first and second dimensions of the input tensor.

The argument offset determines where diagonals are taken from input tensor:

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

  • input (Variable) – The input tensor. Must be at least 2-dimensional. The input data type should be float32, float64, int32, int64.

  • offset (int, optional) – Which diagonals in input tensor will be taken. Default: 0 (main diagonals).

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

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

  • name (str, optional) – Normally there is no need for user to set this property. For more information, please refer to Name. Default: None.


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

Return type



import paddle.fluid as fluid
import paddle.fluid.dygraph as dg
import numpy as np

case1 = np.random.randn(2, 3).astype('float32')
case2 = np.random.randn(3, 10, 10).astype('float32')
case3 = np.random.randn(3, 10, 5, 10).astype('float32')

with dg.guard():
    case1 = dg.to_variable(case1)
    case2 = dg.to_variable(case2)
    case3 = dg.to_variable(case3)
    data1 = fluid.layers.trace(case1) # data1.shape = [1]
    data2 = fluid.layers.trace(case2, offset=1, dim1=1, dim2=2) # data2.shape = [3]
    data3 = fluid.layers.trace(case3, offset=-3, dim1=1, dim2=-1) # data2.shape = [3, 5]