Convert range index to coordinates in tensorflow

Question

I want to convert this kind of indexing:

a = tf.ones((500, 1000))
a[1:150, 50:700]

Into :

a[idx_x, idx_y]

We can obviously use a while loop to construct the two vectors. But what is the most efficient way to do this in tensorflow?

Answer 1

You can generate the index values as:

idx_x, idx_y = tf.meshgrid(tf.range(1, 150), tf.range(50, 700), indexing='ij')

Then you would use those as:

b = tf.gather_nd(a, tf.stack([idx_x, idx_y], axis=-1))