TensorFlow getting elements of every row for specific columns
<p>If <code>A</code> is a TensorFlow variable like so </p>

<pre><code>A = tf.Variable([[1, 2], [3, 4]])
</code></pre>

<p>and <code>index</code> is another variable </p>

<pre><code>index = tf.Variable([0, 1])
</code></pre>

<p>I want to use this index to select columns in each row. In this case, item 0 from first row and item 1 from second row.</p>

<p>If A was a Numpy array then to get the columns of corresponding rows mentioned in index we can do </p>

<pre><code>x = A[np.arange(A.shape[0]), index]
</code></pre>

<p>and the result would be </p>

<pre><code>[1, 4]
</code></pre>

<p>What is the TensorFlow equivalent operation/operations for this? I know TensorFlow doesn't support many indexing operations. What would be the work around if it cannot be done directly?</p>

<p>After dabbling around for quite a while. I found two functions that could be useful.</p>

<p>One is <code>tf.gather_nd()</code> which might be useful if you can produce a tensor 
of the form <code>[[0, 0], [1, 1]]</code> and thereby you could do </p>

<p><code>index = tf.constant([[0, 0], [1, 1]])</code></p>

<p><code>tf.gather_nd(A, index)</code></p>

<p>If you are unable to produce a vector of the form <code>[[0, 0], [1, 1]]</code>(I couldn't produce this as the number of rows in my case was dependent on a placeholder) for some reason then the work around I found is to use the <code>tf.py_func()</code>. Here is an example code on how this can be done </p>

<pre><code>import tensorflow as tf 
import numpy as np 

def index_along_every_row(array, index):
    N, _ = array.shape 
    return array[np.arange(N), index]

a = tf.Variable([[1, 2], [3, 4]], dtype=tf.int32)
index = tf.Variable([0, 1], dtype=tf.int32)
a_slice_op = tf.py_func(index_along_every_row, [a, index], [tf.int32])[0]
session = tf.InteractiveSession()

a.initializer.run()
index.initializer.run()
a_slice = a_slice_op.eval() 
</code></pre>

<p><code>a_slice</code> will be a numpy array <code>[1, 4]</code></p>

python numpy tensorflow