Get the diagonal of a matrix in TensorFlow

Question

Is there a way to extract the diagonal of a square matrix in TensorFlow? That is, for a matrix like this:

[
 [0, 1, 2],
 [3, 4, 5],
 [6, 7, 8]
]

I want to fetch the elements: [0, 4, 8]

In numpy, this is pretty straight-forward via np.diag:

In TensorFlow, there is a diag function, but it only forms a new matrix with the elements specified in the argument on the diagonal, which is not what I want.

I could imagine how this could be done via striding... but I don't see striding for tensors in TensorFlow.

Answer 1

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)

Answer 2

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.

Answer 3

Use the tf.diag_part()

with tf.Session() as sess:
    x = tf.ones(shape=[3, 3])
    x_diag = tf.diag_part(x)
    print(sess.run(x_diag ))