trace

trace(x: array, /, *, offset: int = 0) → array

Returns the sum along the specified diagonals of a matrix (or a stack of matrices) x.

Parameters:

• x (array) – input array having shape (..., M, N) and whose innermost two dimensions form MxN matrices. Should have a numeric data type.

• offset (int) –

  offset specifying the off-diagonal relative to the main diagonal.

  • offset = 0: the main diagonal.

  • offset > 0: off-diagonal above the main diagonal.

  • offset < 0: off-diagonal below the main diagonal.

  Default: 0.

Returns:

out (array) – an array containing the traces and whose shape is determined by removing the last two dimensions and storing the traces in the last array dimension. For example, if x has rank k and shape (I, J, K, ..., L, M, N), then an output array has rank k-2 and shape (I, J, K, ..., L) where

out[i, j, k, ..., l] = trace(a[i, j, k, ..., l, :, :])

The returned array must have the same data type as x.