tf.linalg.tensor_diag_part

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

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

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

Предположим, что тензор имеет размерности , тогда выходной тензор имеет ранг с размерностями , где:

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

Например:

# 'input' is [[1, 0, 0, 0]
              [0, 2, 0, 0]
              [0, 0, 3, 0]
              [0, 0, 0, 4]]

tf.diag_part(input) ==> [1, 2, 3, 4]