how to map tensor to it's indices in tensorflow

Question

Suppose I have a 2D tensor with shape (size, size), and I want to get 2 new tensors that containing the original tensors row index and column index. So if size is 2, I want to get [[0, 0], [1, 1]] and [[0, 1], [0, 1]]

What's tricky is that size is another tensor whose value can only be known when running the graph in a tensorflow Session. How can I do this in tensorflow?

Answer 1

Seems like you are looking for tf.meshgrid.

Here's an example:

shape = tf.shape(matrix)
R, C = tf.meshgrid(tf.range(shape[0]), tf.range(shape[1]), indexing='ij')

matrix is your 2D tensor, R and C contain your row and column indices, respectively. Note that this can be slightly simplified if your matrix is square (only one tf.range).