Range-based indexing based on another tensor

rdesh26: <p>Suppose I have a tensor <strong>x</strong> of shape (B, U, T), and another tensor <strong>t</strong> of shape (B, U, 2). The last dimension in <strong>t</strong> denotes <em>start</em> and <em>end</em> indices for the corresponding T dimension of <strong>x</strong>. I want to create a flattened tensor <strong>y</strong> which contains only the “selected” values of <strong>x</strong>, where the selection is based on the ranges given in <strong>t</strong>. Assume that I know the total size of <strong>y</strong> in advance (this is easy to compute as <code>y_size = (t[:, :, 1] - t[:, :, 0]).sum()</code>).</p>
<p>A naive loop based implement would look something like this:</p>
<pre><code class="lang-python">offset = 0
for b in range(B):
    for u in range(U):
        start, end = t[b, u]
        y[offset : offset + end - start] = x[b, u, start:end]
        offset += end-start
</code></pre>
<p>Of course, this is slow since the loop runs B*U times. I am trying to come up with a more efficient way of doing this, perhaps using <code>index_select()</code>, but I am stuck.</p>
<p>I would appreciate any help if someone can think of a faster solution.</p>

rdesh26: <p>I was able to come up with the following mask based approach, but perhaps there is a more efficient way.</p>
<pre><code class="lang-python">import torch.nn.functional as F

indices = torch.zeros_like(x)
indices += F.one_hot(t[..., 0], num_classes=T)
indices -= F.one_hot(t[..., 1], num_classes=T)
indices = torch.cumsum(indices, dim=-1).bool()

y = torch.masked_select(x, indices)
</code></pre>

KFrank: <p>Hi Desh!</p>
<aside class="quote no-group" data-username="rdesh26" data-post="2" data-topic="158666" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/rdesh26/48/34338_2.png" class="avatar"> rdesh26:</div>
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">indices = torch.zeros_like(x)
indices += F.one_hot(t[..., 0], num_classes=T)
indices -= F.one_hot(t[..., 1], num_classes=T)
indices = torch.cumsum(indices, dim=-1).bool()
</code></pre>
</blockquote>
</aside>
<p>You could also compute <code>indices</code> with a logical expression:</p>
<pre data-code-wrap="python"><code class="lang-python">indices = torch.logical_and (torch.arange (T) &gt;= t[:, :, 0:1], torch.arange (T) &lt; 
t[:, :, 1:2])
</code></pre>
<p>To my eye, this version is a little more readable.  (Because of broadcasting,<br>
<code>indices</code> has the same shape as <code>x</code>.)</p>
<p>I don’t know whether this would have any efficiency benefits.</p>
<p>Best.</p>
<p>K. Frank</p>

rdesh26: <p>Thanks for your suggestion. It is indeed more readable than using one-hot and cumsum. I will try both methods and see if there it makes a difference in the throughput.</p>
<p><strong>Update:</strong> After trying out both methods, I found that the method suggested by <a class="mention" href="/u/kfrank">@KFrank</a> is ~10% slower overall compared to the cumsum. I think this may be because it requires T point-wise comparisons, which is slower than just setting 2 values and doing cumsum on a single tensor.</p>
<p>I also found <a href="https://discuss.pytorch.org/t/set-values-of-1d-tensor-using-list-of-start-end-indices/133458">this related post</a>, where <a class="mention" href="/u/ptrblck">@ptrblck</a> suggests a similar approach.</p>