Get the diagonal of a matrix in TensorFlow

Currently it is possible to extract diagonal elements with tf.diag_part. Here is their example:

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

---

Old answer (when diag_part) was not available (still relevant if you want to achieve something that is not available now):

After looking though the math operations and tensor transformations, it does not look like such operation exists. Even if you can extract this data with matrix multiplications it would not be efficient (get diagonal is O(n)).

You have three approaches, starting with easy to hard.

1. Evaluate the tensor, extract diagonal with numpy, build a variable with TF
2. Use tf.pack in a way Anurag suggested (also extract the value 3 using tf.shape
3. Write your own op in C++, rebuild TF and use it natively.

with tensorflow 0.8 its possible to extract the diagonal elements with tf.diag_part() (see documentation)

UPDATE

for tensorflow >= r1.12 its tf.linalg.tensor_diag_part (see documentation)