title: How to gather the logits according to a 2D index?
body: <p>Assume the sequence <code>[[2, 1, 4], [3, 4, 2]]</code> is generated by a pre-trained LSTM. It's dimension is <code>(2*3)</code> meaning <code>batch-size = 2</code> and <code>3 time steps</code> in each sample.</p>
  
  <p>Then for example, there are <code>5 features</code> in total so the logits may be: </p>
  
  <pre><code>[[[0.2, 0.2, 0.1, 0.1, 0.2], 
    [0.3, 0.2, 0.1, 0.1, 0.1],
    [0.1, 0.2, 0.1, 0.1, 0.3]],
   [[0.2, 0.2, 0.1, 0.1, 0.2],
    [0.2, 0.2, 0.1, 0.1, 0.2],
    [0.2, 0.2, 0.1, 0.1, 0.2]]]
  </code></pre>
  
  <p>I want to use the sequence as the index to get the corresponding probabilities from the logits for each sample and each time step. Regarding the example above, the final result I want to get is </p>
  
  <pre><code>[[0.1, 0.2, 0.3],[0.1, 0.2, 0.1]]
  </code></pre>
  
  <p>I knew that I probably need tf.stack() but I'm confused about how to handle the dimension. Appreciate for any help!</p>