How to use tf.gather similar to the numpy slicing-CodePudding

I have a list of indices representing the rows and columns that I'd like to access.

r= np.array([0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3])
c = np.array([0, 0, 1, 1, 4, 4, 1, 1, 3, 3, 5, 5])

When I use them as [r,c], I am able to obtain the corresponding elements in a numpy array.

arr = np.array([[517.0, 1876.0, 4716.0, 2725.0, 2138.0, 2213.0],
                [517.0, 1876.0, 4716.0, 2725.0, 2138.0, 2213.0],
                [3745.0, 3103.0, 885.0, 3482.0, 4196.0, 1802.0], 
                [3745.0, 3103.0, 885.0, 3482.0, 4196.0, 1802.0],
                [1548.0, 610.0, 3936.0, 905.0, 3791.0, 3657.0], 
                [1548.0, 610.0, 3936.0, 905.0, 3791.0, 3657.0], 
                [971.0, 573.0, 4756.0, 2137.0, 1407.0, 4388.0], 
                [971.0, 573.0, 4756.0, 2137.0, 1407.0, 4388.0], 
                [2786.0, 4769.0, 3391.0, 940.0, 2188.0, 1823.0], 
                [2786.0, 4769.0, 3391.0, 940.0, 2188.0, 1823.0], 
                [3225.0, 3262.0, 3444.0, 783.0, 3931.0, 1546.0], 
                [3225.0, 3262.0, 3444.0, 783.0, 3931.0, 1546.0]])
print(mst_array[r, c])
[ 517.  517. 1876. 1876. 2138. 2138. 3103. 3103. 3482. 3482. 1802. 1802.]

I would like to perform the same operation in tf and am using tf.gather.

https://www.tensorflow.org/api_docs/python/tf/gather

t_r = tf.convert_to_tensor(r, dtype= tf.int32)
t_c =  tf.convert_to_tensor(c, dtype= tf.int32)
t_arr = tf.convert_to_tensor(arr, dtype= tf.int32)
print(t_r, "\n", t_c, "\n")
tf.Tensor([0 1 0 1 0 1 2 3 2 3 2 3], shape=(12,), dtype=int32) 
tf.Tensor([0 0 1 1 4 4 1 1 3 3 5 5], shape=(12,), dtype=int32)

I have tried different commands, but could not manage to get the same results as I do in numpy.

tf.gather(t_arr, [t_r, t_c], axis=-1)
tf.gather(t_arr, [t_r, t_c], axis=-1).shape_as_list()

CodePudding user response：

Try tf.stack with tf.gather_nd:

z = tf.gather_nd(t_arr, tf.stack([t_r, t_c], axis=1))

tf.Tensor([ 517  517 1876 1876 2138 2138 3103 3103 3482 3482 1802 1802], shape=(12,), dtype=int32)