title: `torch.gather` without unbroadcasting
body: <p>I have some batched input <code>x</code> of shape <code>[batch, time, feature]</code>, and some batched indices <code>i</code> of shape <code>[batch, new_time]</code> which I want to gather into the time dim of <code>x</code>. As output of this operation I want a tensor <code>y</code> of shape <code>[batch, new_time, feature]</code> with values like this:</p>
  <pre><code>y[b, t', f] = x[b, i[b, t'], f]
  </code></pre>
  <p>In Tensorflow, I can accomplish this by using the <a href="https://www.tensorflow.org/api_docs/python/tf/gather" rel="nofollow noreferrer"><code>batch_dims: int</code> argument of <code>tf.gather</code></a>: <code>y = tf.gather(x, i, axis=1, batch_dims=1)</code>.</p>
  <p>In PyTorch, I can think of some functions which do similar things:</p>
  <ol>
  <li><p><a href="https://pytorch.org/docs/stable/generated/torch.gather.html" rel="nofollow noreferrer"><code>torch.gather</code></a> of course, but this does not have an argument similar to Tensorflow's <code>batch_dims</code>. The output of <code>torch.gather</code> will always have the same shape as the indices. So I would need to unbroadcast the <code>feature</code> dim into <code>i</code> before passing it to <code>torch.gather</code>.</p>
  </li>
  <li><p><a href="https://pytorch.org/docs/stable/generated/torch.index_select.html" rel="nofollow noreferrer"><code>torch.index_select</code></a>, but here, the indices must be one-dimensional. So to make it work I would need to unbroadcast <code>x</code> to add a &quot;<code>batch * new_time</code>&quot; dim, and then after <code>torch.index_select</code> reshape the output.</p>
  </li>
  <li><p><a href="https://pytorch.org/docs/stable/generated/torch.nn.functional.embedding.html" rel="nofollow noreferrer"><code>torch.nn.functional.embedding</code></a>. Here, the embedding matrices would correspond to <code>x</code>. But this embedding function does not support the weights to be batched, so I run into the same issue as for <code>torch.index_select</code> (looking at the code, <a href="https://github.com/pytorch/pytorch/blob/b85568a54a9c60986235ad1e0cc5dffc71b9d5b1/aten/src/ATen/native/Embedding.cpp#L14" rel="nofollow noreferrer"><code>tf.embedding</code> uses <code>torch.index_select</code> under the hood</a>).</p>
  </li>
  </ol>
  <p><strong>Is it possible to accomplish such gather operation without relying on unbroadcasting which is inefficient for large dims?</strong></p>