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].

Пример:

# '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]