tf.keras.ops.trace

Возвращает сумму по диагоналям тензора.

tf.keras.ops.trace(
    x, offset=0, axis1=0, axis2=1
)

Если x является 2-мерным, возвращается сумма по диагонали с заданным смещением, т.е. сумма элементов x[i, i+offset] для всех i.

Если у a есть более двух измерений, то для определения 2-мерных подмассивов, чьи следы возвращаются, используются оси, указанные в axis1 и axis2.

Форма результирующего тензора такая же, как у x, с удаленными axis1 и axis2.