tf.keras.ops.diagonal

Возвращает указанные диагонали.

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

Если x — двумерный массив, возвращает диагональ x с указанным смещением, то есть набор элементов вида x[i, i+offset].

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

Форму результирующего массива можно определить, удалив axis1 и axis2 и добавив справа индекс, равный размеру результирующих диагоналей.

Примеры:

from keras.src import ops
x = ops.arange(4).reshape((2, 2))
x
array([[0, 1],
       [2, 3]])
x.diagonal()
array([0, 3])
x.diagonal(1)
array([1])

x = ops.arange(8).reshape((2, 2, 2))
x
array([[[0, 1],
        [2, 3]],
       [[4, 5],
        [6, 7]]])
x.diagonal(0, 0, 1)
array([[0, 6],
       [1, 7]])