theaNO theaNO -4 years ago 157
Python Question

Get the diagonal of a matrix in TensorFlow

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 Source

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

Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download