tf.linalg.tensor_diag_part

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

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

Данная операция возвращает тензор с диагональной частью тензора. Диагональная часть вычисляется следующим образом:

Предположим, что 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])