Tensorflow: how to obtain the diagonals of multiple matrices (batched mode)

Question

I have a tensor which is [100 X 16 X 16]. I want to get the diagonal elements of this tensor to get a tensor of shape [100 X 16]. I tried the following:

#sum_cov is [100 X 16 X 16] and diagonal_elements is expected to be [100 X 16].

diagonal_elements = tf.diag_part(sum_cov)

But, I get the following error:

Input must have even rank <= 6, input rank is 3 for 'DiagPart'

Can someone please tell me how to achieve this?

Answer 1

https://www.tensorflow.org/api_docs/python/tf/diag_part https://www.tensorflow.org/api_docs/python/tf/matrix_diag_part

dm0_ is right. You want tf.matrix_diag_part.

tf.diag_part computes the tensor diagonal, where if the input tensor has shape (D1, ..., Dk, D1, ..., Dk) then the output tensor has shape (D1, ..., Dk) and is such that

tf.diag_part(input)[i1, ..., ik] = input[i1, ..., ik, i1, ..., ik]

This is why you're getting an error. In order for the above preconditions to hold, the input tensor must have even rank.

tf.matrix_diag_part, on the other hand, treats the input tensor as a batch of 2-dimensional matrices, and computes the diagonal of each. So if the input tensor has shape (I, J, K, ..., M, N), the output tensor will have shape (I, J, K, ..., min(M, N)) and be such that

tf.matrix_diag_part(input)[i, j, k, ..., m] = input[i, j, k, ..., m, m]

The two functions are identical for rank 2 tensors, but for anything higher than that they're very different beasts.