title: How to select rows from a 3-D Tensor in TensorFlow?
body: <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>