How to select rows from a 3-D Tensor in TensorFlow?
<p>I have a tensor <code>logits</code> with the dimensions <code>[batch_size, num_rows, num_coordinates]</code> (i.e. each logit in the batch is a matrix). In my case batch size is 2, there's 4 rows and 4 coordinates.</p>

<pre><code>logits = tf.constant([[[10.0, 10.0, 20.0, 20.0],
                      [11.0, 10.0, 10.0, 30.0],
                      [12.0, 10.0, 10.0, 20.0],
                      [13.0, 10.0, 10.0, 20.0]],
                     [[14.0, 11.0, 21.0, 31.0],
                      [15.0, 11.0, 11.0, 21.0],
                      [16.0, 11.0, 11.0, 21.0],
                      [17.0, 11.0, 11.0, 21.0]]])
</code></pre>

<p>I want to select the first and second row of the first batch and the second and fourth row of the second batch.</p>

<pre><code>indices = tf.constant([[0, 1], [1, 3]])
</code></pre>

<p>So the desired output would be</p>

<pre><code>logits = tf.constant([[[10.0, 10.0, 20.0, 20.0],
                      [11.0, 10.0, 10.0, 30.0]],
                     [[15.0, 11.0, 11.0, 21.0],
                      [17.0, 11.0, 11.0, 21.0]]])
</code></pre>

<p>How do I do this using TensorFlow? I tried using <code>tf.gather(logits, indices)</code> but it did not return what I expected. Thanks!</p>