tf.linalg.tensor_diag_part

Возвращает диагональную часть тензора.

tf.linalg.tensor_diag_part(
    input, name=None
)

Эта операция возвращает тензор, содержащий diagonal часть input. diagonal часть вычисляется следующим образом:

Предположим, что input имеет размеры [D1,..., Dk, D1,..., Dk], тогда выходные данные — это тензор ранга k с размерами [D1,..., Dk], где:

diagonal[i1,..., ik] = input[i1, ..., ik, i1,..., ik].

Для тензора ранга 2 функции linalg.diag_part и linalg.tensor_diag_part дают одинаковый результат. Для ранга 3 и выше linalg.diag_part извлекает диагональ каждой самой внутренней матрицы в тензоре. Ниже приведен пример, в котором результаты различаются.

x = [[[[1111,1112],[1121,1122]],
      [[1211,1212],[1221,1222]]],
     [[[2111, 2112], [2121, 2122]],
      [[2211, 2212], [2221, 2222]]]
     ]
tf.linalg.tensor_diag_part(x)
<tf.Tensor: shape=(2, 2), dtype=int32, numpy=
array([[1111, 1212],
       [2121, 2222]], dtype=int32)>
tf.linalg.diag_part(x).shape
TensorShape([2, 2, 2])