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 such a matrix:

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

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

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

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

I could imagine how this could be done with a step ... but I do not see the tensors moving in TensorFlow.

Answer 1

with a tensor flow of 0.8, you can extract diagonal elements using tf.diag_part() (see documentation )

UPDATE

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

Answer 2

This is probably a workaround, but works.

 >> sess = tensorflow.InteractiveSession() >> x = tensorflow.Variable([[1,2,3],[4,5,6],[7,8,9]]) >> x.initializer.run() >> z = tensorflow.pack([x[i,i] for i in range(3)]) >> z.eval() array([1, 5, 9], dtype=int32) 

Answer 3

Currently, you can extract diagonal elements using tf.diag_part . Here is an example of them:

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

---

The old answer (when diag_part) is unavailable (still relevant if you want to achieve something that isn’t right now):

After looking at the mathematical operations and tensor transformations , it seems that such an operation does not exist. Even if you can extract this data using matrix multiplication, it will be inefficient (get the diagonal O(n) ).

You have three approaches, starting with simple and complex.

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

Answer 4

Use 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 )) 

Answer 5

Use the gather operation.

 x = tensorflow.Variable([[1,2,3],[4,5,6],[7,8,9]]) x_flat = tf.reshape(x, [-1]) # flatten the matrix x_diag = tf.gather(x, [0, 3, 6]) 

Answer 6

Depending on the context, a mask can be a good way to “undo” the diagonal elements of a matrix, especially if you plan to reduce it:

 mask = tf.diag(tf.ones([n])) y = tf.mul(mask,y) cost = -tf.reduce_sum(y)