Efficient way to do a strided sum and reduce a nxm tensor to a n vector

btnadiga: <p>Typically m &lt;&lt; n. But for example, say with n=3, m=2, I want to reduce a 3x2 tensor to a 3 vector in this fashion:<br>
I want to write the 3x2 tensor in this fashion</p>
<p>00  01   xx<br>
xx  10  11<br>
21   xx   20</p>
<p>where I’m only indicating the indices, xx indicates elements that don’t exist and then take a column sum to obtain a 3 vector<br>
[00 + 21, 01 + 10, 11 + 20]</p>
<p>Or without the wrap-around, write it as<br>
00  01  xx    xx<br>
xx  10   11    xx<br>
xx   xx   20  21<br>
and again sum the columns to obtain<br>
[00, 01 + 10, 11 + 20, 21][:-1]</p>
<p>seems almost like a conv1d-like operation or torch.diag/torch.diagonal operation<br>
torch.sum( [ torch.roll( torch.diag( y[ :, i ] ), shifts=i, dims=1 ) for i in range( len( y[ 0 ] ) ], axis=0 ) ?</p>
<p>Thanks</p>
<p>ps: When I had spaces instead of xx in the original post, the spaces were eaten up and the displayed version was not how I typed it</p>

KFrank: <p>Hi btnadiga!</p>
<aside class="quote no-group" data-username="btnadiga" data-post="1" data-topic="161939" 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/btnadiga/48/51773_2.png" class="avatar"> btnadiga:</div>
<blockquote>
<p>I want to write the 3x2 tensor in this fashion</p>
<p>00  01   xx<br>
xx  10  11<br>
21   xx   20</p>
<p>where I’m only indicating the indices, xx indicates elements that don’t exist and then take a column sum to obtain a 3 vector<br>
[00 + 21, 01 + 10, 11 + 20]</p>
</blockquote>
</aside>
<p>One approach would be to append (<code>cat()</code>) a column of zeros to your<br>
original tensor, use torch tensor indexing with appropriately constructed<br>
indices to obtain your “shifted” tensor, and then perform the sum:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.12.0'
&gt;&gt;&gt; t = 10 * torch.arange (10).unsqueeze (1) + torch.arange (2)
&gt;&gt;&gt; t
tensor([[ 0,  1],
        [10, 11],
        [20, 21],
        [30, 31],
        [40, 41],
        [50, 51],
        [60, 61],
        [70, 71],
        [80, 81],
        [90, 91]])
&gt;&gt;&gt; t0 = torch.cat ((t, torch.zeros (10, 1, dtype = torch.long)), 1)
&gt;&gt;&gt; t0
tensor([[ 0,  1,  0],
        [10, 11,  0],
        [20, 21,  0],
        [30, 31,  0],
        [40, 41,  0],
        [50, 51,  0],
        [60, 61,  0],
        [70, 71,  0],
        [80, 81,  0],
        [90, 91,  0]])
&gt;&gt;&gt; t0[torch.arange (10).unsqueeze (1).expand (10, 3), (torch.arange (30).view (10, 3) - torch.arange (10).unsqueeze (1)) % 3]
tensor([[ 0,  1,  0],
        [ 0, 10, 11],
        [21,  0, 20],
        [30, 31,  0],
        [ 0, 40, 41],
        [51,  0, 50],
        [60, 61,  0],
        [ 0, 70, 71],
        [81,  0, 80],
        [90, 91,  0]])
&gt;&gt;&gt; t0[torch.arange (10).unsqueeze (1).expand (10, 3), (torch.arange (30).view (10, 3) - torch.arange (10).unsqueeze (1)) % 3].sum (0)
tensor([333, 304, 273])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

btnadiga: <p><a class="mention" href="/u/kfrank">@KFrank</a>, coming from numpy my solution was</p>
<pre><code class="lang-auto">import torch
n, m = 3, 2
y2d = torch.arange(6.).reshape(n, m)
print(y2d)
tmp = torch.sum ( torch.stack( [ torch.roll( torch.diag( y2d[:, i] ),
                                             shifts=i, dims=1 )
                                 for i in range(m) ] ), axis=(0, 1) )
print(tmp)
</code></pre>
<p>Any idea which is better in the sense of speed?<br>
Thanks!<br>
Balu</p>

KFrank: <p>Hi Balu!</p>
<aside class="quote no-group" data-username="btnadiga" data-post="3" data-topic="161939" 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/btnadiga/48/51773_2.png" class="avatar"> btnadiga:</div>
<blockquote>
<p>Any idea which is better in the sense of speed?</p>
</blockquote>
</aside>
<p>If efficiency actually matters, I would suggest that you time both versions.</p>
<p>Using loop-free pytorch tensor operations will typically outperform (often<br>
by a lot) an equivalent for-loop / list-comprehension computation.  But<br>
that’s just a rule of thumb that may or may not apply to your use case.</p>
<p>Best.</p>
<p>K. Frank</p>