How to extract rows and columns from a 3D array in Tensorflow

Question

I wanted to do the following indexing operation on a TensorFlow tensor. What should be the equivalent operations in TensorFlow to get b and c as output? Although tf.gather_nd documentation has several examples but I could not generate equivalent indices tensor to get these results.

import tensorflow as tf
import numpy as np

a=np.arange(18).reshape((2,3,3))

idx=[2,0,1] #it can be any validing re-ordering index list

#These are the two numpy operations that I want to do in Tensorflow
b=a[:,idx,:]
c=a[:,:,idx] 

# TensorFlow operations

aT=tf.constant(a)
idxT=tf.constant(idx)

# what should be these two indices  
idx1T=tf.reshape(idxT, (3,1)) 
idx2T=tf.reshape(idxT, (1,1,3))

bT=tf.gather_nd(aT, idx1T ) #does not work
cT=tf.gather_nd(aT, idx2T)  #does not work

with tf.Session() as sess:
    b1,c1=sess.run([bT,cT])

print(np.allclose(b,b1))
print(np.allclose(c,c1))

I am not restricted to tf.gather_nd Any other suggestion to achieve the same operations on GPU will be helpful.

Edit: I have updated the question for a typo:

old statement: c=a[:,idx],

New statement: c=a[:,:,idx] What I wanted to achieve was re-ordering of columns as well.

javidcf · Accepted Answer · 2018-10-18 16:24:48Z

2

That can be done with tf.gather, using the axis parameter:

import tensorflow as tf
import numpy as np

a = np.arange(18).reshape((2,3,3))
idx = [2,0,1]
b = a[:, idx, :]
c = a[:, :, idx]

aT = tf.constant(a)
idxT = tf.constant(idx)
bT = tf.gather(aT, idxT, axis=1)
cT = tf.gather(aT, idxT, axis=2)

with tf.Session() as sess:
    b1, c1=sess.run([bT, cT])

print(np.allclose(b, b1))
print(np.allclose(c, c1))

Output:

True
True

answered Oct 18, 2018 at 16:24

javidcf

59.9k7 gold badges87 silver badges134 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Hemant Over a year ago

Thanks @jdehesa, yes I meant c=a[:,:,idx]. Updated it in question.

Collectives™ on Stack Overflow

How to extract rows and columns from a 3D array in Tensorflow

Edit: I have updated the question for a typo:

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

Edit: I have updated the question for a typo:

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related