CODEWITHSUNDEEP
pythontensorflow
Dims
Dims

Reputation: 51039

How to select different columns from each row in tensor in tensorflow?

This question is the same as this one, but for Tensorflow.

Suppose I have 2D tensor of "rows" and want to select i-th element from each row and compose result column of these elements, having i-s in selector tensor, like below

import tensorflow as tf
import numpy as np

rows = tf.constant(np.arange(10*3).reshape(10,3), dtype=tf.float64)
# gives
# array([[ 0,  1,  2],
#        [ 3,  4,  5],
#        [ 6,  7,  8],
#        [ 9, 10, 11],
#        [12, 13, 14],
#        [15, 16, 17],
#        [18, 19, 20],
#        [21, 22, 23],
#        [24, 25, 26],
#        [27, 28, 29]])


selector = tf.get_variable("selector", [10,1], dtype=tf.int8, initializer=tf.constant([[0], [1], [0], [2], [1], [0], [0], [2], [2], [1]]))

result_of_selection = ...

# should be
# array([[ 0],
#        [ 4],
#        [ 6],
#        [11],
#        [13],
#        [15],
#        [18],
#        [23],
#        [26],
#        [28]])

How would I do this?

UPDATE

I wrote this way (thanks @Psidom) 

import tensorflow as tf
import numpy as np

rows = tf.constant(np.arange(10*3).reshape(10,3), dtype=tf.float64)

# selector = tf.get_variable("selector", dtype=tf.int32, initializer=tf.constant([0, 1, 0, 2, 1, 0, 0, 2, 2, 1], dtype=tf.int32))
# selector = tf.expand_dims(selector, axis=1)
selector = tf.get_variable("selector", dtype=tf.int32, initializer=tf.constant([[0], [1], [0], [2], [1], [0], [0], [2], [2], [1]], dtype=tf.int32))

ordinals = tf.reshape(tf.range(rows.shape[0]), (-1,1))

#idx = tf.concat([selector, ordinals], axis=1)
idx = tf.stack([selector, ordinals], axis=-1)

result = tf.gather_nd(rows, idx)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    rows_value, result_value = sess.run([rows, result])
    print("rows_value: " + str(rows_value))
    print("selector_value: " + str(result_value))

and it gave

rows_value: [[  0.   1.   2.]
 [  3.   4.   5.]
 [  6.   7.   8.]
 [  9.  10.  11.]
 [ 12.  13.  14.]
 [ 15.  16.  17.]
 [ 18.  19.  20.]
 [ 21.  22.  23.]
 [ 24.  25.  26.]
 [ 27.  28.  29.]]
selector_value: [[ 0.]
 [ 4.]
 [ 2.]
 [ 0.]
 [ 0.]
 [ 0.]
 [ 0.]
 [ 0.]
 [ 0.]
 [ 0.]]

I.e. incorrect.

UPDATE 2

Fixed line 

idx = tf.stack([ordinals, selector], axis=-1)

was incorrect order.

Upvotes: 0

Views: 2679

Answers (1)

akuiper
akuiper

Reputation: 214957

One way to do this is explicitly construct the indices by stacking the row indices which can be created using tf.range with the selector, and then use tf.gather_nd to collect the items:

rows = tf.constant(np.arange(10*3).reshape(10,3), dtype=tf.float64)
selector = tf.constant([[0], [1], [0], [2], [1], [0], [0], [2], [2], [1]])

idx = tf.stack([tf.reshape(tf.range(rows.shape[0]), (-1,1)), selector], axis=-1)

with tf.Session() as sess:
    print(sess.run(tf.gather_nd(rows, idx)))

#[[  0.]
# [  4.]
# [  6.]
# [ 11.]
# [ 13.]
# [ 15.]
# [ 18.]
# [ 23.]
# [ 26.]
# [ 28.]]

Here idx is the actual index of all the elements in the original tensor:

with tf.Session() as sess:
    print(idx.eval())
#[[[0 0]]

# [[1 1]]

# [[2 0]]

# [[3 2]]

# [[4 1]]

# [[5 0]]

# [[6 0]]

# [[7 2]]

# [[8 2]]

# [[9 1]]]

Edit: selector as a variable:

selector = tf.get_variable("selector", dtype=tf.int32, initializer=tf.constant([[0], [1], [0], [2], [1], [0], [0], [2], [2], [1]]))
idx = tf.stack([tf.reshape(tf.range(rows.shape[0]), (-1,1)), selector], axis=-1)

with tf.Session() as sess:
    tf.global_variables_initializer().run()
    print(sess.run(tf.gather_nd(rows, idx)))

#[[  0.]
# [  4.]
# [  6.]
# [ 11.]
# [ 13.]
# [ 15.]
# [ 18.]
# [ 23.]
# [ 26.]
# [ 28.]]

Upvotes: 3

Related Questions