Generating a 2D matrix with 1s at particular positions

souryadey: <p>I have a 1D tensor <code>s</code> with elements in ascending order. I want to generate a 2D tensor <code>a</code> of shape <code>(len(s)-1, s[-1])</code> such that its elements <code>[i, s[i]:s[i+1]]</code> are <code>1</code>, all other elements being <code>0</code>. How can I do this vectorially without using loops?</p>
<p>Example:</p>
<p>Input: <code>s = torch.tensor([0,2,5,7])</code><br>
Output:</p>
<pre><code class="lang-python">a = torch.tensor([
    [1,1,0,0,0,0,0],
    [0,0,1,1,1,0,0],
    [0,0,0,0,0,1,1]
])
</code></pre>
<p>Loop implementation:</p>
<pre><code class="lang-python">a = torch.zeros(len(s)-1,s[-1])
for i in range(len(a)):
    a[i, s[i]:s[i+1]] = 1
</code></pre>
<p>This is part of training a neural network, so I’m looking for a vector implementation.</p>

souryadey: <p>In case it helps, this is actually part of a bigger problem where I am trying to sum unevenly-sized rows of a 2D tensor.</p>
<p>As an example:</p>
<pre><code class="lang-python">inp = torch.tensor([
    [0,1],
    [2,3],
    [4,5],
    [6,7],
    [8,9],
    [10,11],
    [12,13]
])
s = torch.tensor([0,2,5,7])
</code></pre>
<p>I want my output to be of 3 rows, 0th row should sum rows <code>0:2</code> of <code>inp</code>, 1’th’ row should sum rows <code>2:5</code> of <code>inp</code>, and 2’th’ row should sum rows <code>5:7</code> of <code>inp</code>. So, the output should be:</p>
<pre><code class="lang-python">out = torch.tensor([
    [2,4],
    [18,21],
    [22,24]
])
</code></pre>
<p>I realized I could do this if I can generate <code>a</code> as above, then <code>out = a@inp</code>.</p>
<p>I know I can solve the problem using <code>scatter_add</code>, but I don’t know how to construct the <code>index</code> parameter of <code>scatter_add</code> from <code>s</code> without using loops.</p>

rinkujadhav2013: <p>I think you will have to loop over the splits.<br>
But here’s how it can be done:<br>
Convert your s array to an array representing the length of each split. So <code>s = [2-0, 5-2, 7-5] = [2,3,2]</code><br>
Then use <code>split</code>, <code>sum</code>, <code>stack</code> as follows:<br>
<code>torch.stack([i.sum(0) for i in torch.split(inp, [2,3,2])])</code></p>

souryadey: <p>Thanks, yeah that’s definitely a good way of doing it.</p>
<p>But I’m still wondering if there’s a non-loop way of doing it in time complexity O(1). Something like <code>torch.sum(inp, splits=s)</code> that would automatically sum over contiguous chunks of uneven sizes and collect them together.</p>

J_Johnson: <p>Not the prettiest solution, but it does away with loops.</p>
<pre><code class="lang-auto">import torch

s=torch.tensor([0,2,5,7])
w=s.size()[0]-1

sp=torch.stack([torch.zeros(w),s[:-1],s[1:], torch.full((w,), s.max().item())]).T
sp = sp[:,1:]-sp[:,:-1]
sp = sp.int()

z=torch.stack([torch.cat([torch.zeros(sp[i][0].item()), torch.ones(sp[i][1].item()), torch.zeros(sp[i][2].item()) ]) for i in range(sp.size()[0])])
print(z)
</code></pre>
<p>EDIT: That should have been zero on the <code>range(sp.size()[0])</code>. Should work now. Please try again.</p>

J_Johnson: <p>Here is a bigger test on the above method:</p>
<pre><code class="lang-auto">import torch

s=torch.arange(10000)*2
print(s)
</code></pre>
<blockquote>
<p>tensor([    0,     2,     4,  …, 19994, 19996, 19998])</p>
</blockquote>
<p>So the final size should be 19998 on dim 1.</p>
<pre><code class="lang-auto">w=s.size()[0]-1
sp=torch.stack([torch.zeros(w),s[:-1],s[1:], torch.full((w,), s.max().item())]).T
sp = sp[:,1:]-sp[:,:-1]
sp = sp.int()
print(sp)
</code></pre>
<blockquote>
<p>tensor([[    0,     2, 19996],<br>
[    2,     2, 19994],<br>
[    4,     2, 19992],<br>
…,<br>
[19992,     2,     4],<br>
[19994,     2,     2],<br>
[19996,     2,     0]], dtype=torch.int32)</p>
</blockquote>
<pre><code class="lang-auto">z=torch.stack([torch.cat([torch.zeros(sp[i][0].item()), torch.ones(sp[i][1].item()), torch.zeros(sp[i][2].item()) ]) for i in range(sp.size()[0])])
print(z)
</code></pre>
<blockquote>
<p>tensor([[1., 1., 0.,  …, 0., 0., 0.],<br>
[0., 0., 1.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
…,<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 1., 0., 0.],<br>
[0., 0., 0.,  …, 0., 1., 1.]])</p>
</blockquote>
<pre><code class="lang-auto">print(z.size())
</code></pre>
<blockquote>
<p>torch.Size([9999, 19998])</p>
</blockquote>
<p>Total run time was under 1 second on CPU.</p>

souryadey: <p>Thank you, that’s a clever solution! I really appreciate you taking the time to write a test case.</p>
<p>I noticed that it still uses loops, i.e. <code>for i in range(sp.size()[0])</code>. I did a time comparison for 1) <a class="mention" href="/u/rinkujadhav2013">@rinkujadhav2013</a> 's solution, vs 2) this. 1) is a bit faster.</p>
<p>In conclusion, it seems like loops are unavoidable for this problem <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>Appreciate both of your efforts!</p>

rinkujadhav2013: <p>I think loops might be unavoidable cuz your splits are of uneven size. When doing matrix multiplications, dimensions need to align across things being parallelised.</p>

souryadey: <p>Well, torch_scatter is built to handle splits of uneven size vectorially, provided one can provide the correct arguments to it. In my case, constructing the <code>index</code> argument to <code>scatter_add</code> needed a loop, that’s why I was looking at alternatives.</p>

rinkujadhav2013: <p>Torch <code>scatter</code> and <code>gather</code> use loops as far as I can tell.<br>
Check the Cuda kernel here: <a href="https://github.com/pytorch/pytorch/blob/d7847ed23eb8880c7541319aab8864a45e239994/aten/src/ATen/native/cuda/ScatterGatherKernel.cu#L88-L99" class="inline-onebox" rel="noopener nofollow ugc">pytorch/ScatterGatherKernel.cu at d7847ed23eb8880c7541319aab8864a45e239994 · pytorch/pytorch · GitHub</a><br>
or the CPU kernels here:<br>
<a href="https://github.com/pytorch/pytorch/blob/8f5f15a64b7f2efabfc9137ddc1ffb3390375626/aten/src/ATen/native/cpu/ScatterGatherKernel.cpp" class="inline-onebox" rel="noopener nofollow ugc">pytorch/ScatterGatherKernel.cpp at 8f5f15a64b7f2efabfc9137ddc1ffb3390375626 · pytorch/pytorch · GitHub</a></p>

souryadey: <p>Oh, that’s interesting. Probably explains why a torch scatter solution I came up with (with loops to construct the <code>index</code> matrix) took longer than your suggested <code>torch.stack</code> solution.</p>

J_Johnson: <p><a class="mention" href="/u/souryadey">@souryadey</a> when a for loop is contained within square brackets (i.e. <code>[k for k in x]</code>) it becomes a special asynchronous iterator in Python. Asynchronous means each step of k computes separately and doesn’t wait for other iters of k to compute first. So it’s effectively not a “for loop”, though still uses for.</p>

souryadey: <p>My understanding is that something like <code>[k for k in x]</code> is a loop comprehension, which may be faster than a regular loop in certain settings. But it doesn’t give the performance of vector processing. <a href="https://towardsdatascience.com/list-comprehensions-vs-for-loops-it-is-not-what-you-think-34071d4d8207" rel="noopener nofollow ugc">This article</a> has a good discussion.</p>

arojnuck: <p>Try this?</p>
<pre><code class="lang-auto">import numpy as np

a = [0,2,5,7]
a_1 = a[:-1]
a_2 = a[1:]
b = np.max(a)+1 # or however many columns you would like +1
(np.tril(np.ones((b, b)), -1)[a_2] * np.tril(np.ones((b, b)), 0)[::-1, ::-1][a_1])[:, :-1]
</code></pre>
<p>Sorry not super familiar with pytorch yet, but I know there are torch equivalents of all of these functions. I think flipping tensors ([::-1] in np) might be different though.</p>