title: torch.narrow() with tensor of start indices
body: <p>I'm wondering if it's possible to do something in PyTorch...</p>
  <p>Basically, what I'm trying to do is something like <code>torch.narrow()</code>, but where the <code>start</code> value is a tensor of indices.</p>
  <p>Suppose I had the following tensors:</p>
  <pre class="lang-py prettyprint-override"><code>start_indices = [3, 1, 2, ...] # shape: (M,)
  dataset = [
      3, 5, 3, 4, 8, 0, 1,
      3, 9, 7, 2, 7, 3, 7,
      6, 0, 2, 3, 0, 2, 5,
      ...
  ] # shape: (M, N) (here, N = 7, just for demo)
  </code></pre>
  <p>In my use case, <code>M ~ 10^3</code> and <code>N &gt; 10^6</code></p>
  <p>I'm wondering if it's possible to do something like this:</p>
  <pre class="lang-py prettyprint-override"><code>torch.narrow(
      input=dataset,
      dim=1,
      start=start_indices,
      length=4
  )
  </code></pre>
  <p>that returns a &quot;narrowed&quot; version of the original dataset, where the start of each row is selected by its corresponding start index. To complete the example I gave above, the result would look like:</p>
  <pre class="lang-py prettyprint-override"><code>result = [
      4, 8, 0, 1,
      9, 7, 2, 7,
      2, 3, 0, 2,
      ....
  ] # shape: (M, 4)
  </code></pre>
  <p>I realize that depending on the start indices and length, you could have some lines going out of bounds of the original tensor -&gt; I'm fine with receiving an error in this case.</p>
  <p>Right now, I'm currently using a for-loop to do this, but it becomes far too slow as <code>M</code> becomes large. I'm hoping there's a way to do this that takes advantage of data parallelism.</p>
  <p>Any ideas? I thought of using <code>torch.index_select()</code>, but it only selects within rows based off of single indices. I want to select within rows by single indices, but then <em>extend</em> those selections out by some fixed length. In that way, what I'm looking for is kind of a &quot;combination&quot; of <code>torch.narrow()</code> and <code>torch.index_select()</code>.</p>
  <p>Thanks in advance!</p>