"Cumulative" average of every N elements along a tensor dimension

imperialn00b: <p>Hi, I want to compute the average of every <code>N</code> elements in a 1D torch tensor of size <code>M</code>. As an example, if my tensor was <code>tensor([7, 4, 8, 2, 6])</code> and <code>N = 3</code>, I want the output to be <code>tensor([6.3333, 4.6667, 5.3333])</code> - since <code>(7+4+8)/3 = 6.3333</code>, <code>(4+8+2)/3 = 4.6667</code>, and <code>(8+2+6)/3 = 5.3333</code>.</p>
<p>To be precise, my input tensor has dimensions <code>(d_0, d_1, ..., d_{n-1}, d_n)</code>. I take as input two variables - <code>1 &lt;= N &lt;= d_{n-1}</code> and <code>0 &lt;= idx &lt;= d_n</code>. I want as output a tensor with dimensions <code>(d_0, d_1, ..., d_{n-1} - N + 1)</code> where I have averaged over every <code>N</code> elements of index <code>idx</code> of dimension <code>d_n</code> as before. Here’s some code showing how I do it at the moment and the result:</p>
<pre><code class="lang-auto">d0, d1, d2 = 2, 5, 4
N, idx = 3, 2
input = torch.randint(0, 10, (d0, d1, d2)).float()
output = torch.cat([torch.tensor([input[d0idx, i:i+N, idx].mean(dim=-1) for i in range(0, d1 - N + 1)]).unsqueeze(0) for d0idx in range(input.shape[0])], dim=0)
print(input)
&gt;&gt;&gt; tensor([[[0., 0., 3., 7.],
         [5., 1., 2., 4.],
         [2., 8., 7., 5.],
         [4., 0., 0., 2.],
         [2., 8., 5., 3.]],

        [[1., 7., 1., 4.],
         [8., 4., 4., 7.],
         [3., 8., 4., 0.],
         [0., 4., 1., 0.],
         [6., 4., 7., 2.]]])

print(output)
&gt;&gt;&gt; tensor([[4., 3., 4.],
        [3., 3., 4.]])

print(output.shape)
&gt;&gt;&gt; torch.Size([2, 3])
</code></pre>
<p>Ideally I do not want to implement this like I have right now with the nested list comprehension. What is the best way to do this?</p>
<p>Thanks!</p>

KFrank: <p>Hi Imperial!</p>
<aside class="quote no-group" data-username="imperialn00b" data-post="1" data-topic="201553" 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/imperialn00b/48/69550_2.png" class="avatar"> imperialn00b:</div>
<blockquote>
<p>As an example, if my tensor was <code>tensor([7, 4, 8, 2, 6])</code> and <code>N = 3</code>, I want the output to be <code>tensor([6.3333, 4.6667, 5.3333])</code> - since <code>(7+4+8)/3 = 6.3333</code>, <code>(4+8+2)/3 = 4.6667</code>, and <code>(8+2+6)/3 = 5.3333</code>.</p>
</blockquote>
</aside>
<p>Set up a <a href="https://pytorch.org/docs/2.3/generated/torch.nn.Conv1d.html" rel="noopener nofollow ugc">Conv1d</a> that has the desired kernel.  For the example you give, you<br>
would want a <code>kernel_size</code> of <code>3</code> and you would set all three kernel values to<br>
<code>0.33333</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>To be precise, my input tensor has dimensions <code>(d_0, d_1, ..., d_{n-1}, d_n)</code>. I take as input two variables - <code>1 &lt;= N &lt;= d_{n-1}</code> and <code>0 &lt;= idx &lt;= d_n</code>. I want as output a tensor with dimensions <code>(d_0, d_1, ..., d_{n-1} - N + 1)</code> where I have averaged over every <code>N</code> elements of index <code>idx</code> of dimension <code>d_n</code> as before.</p>
</blockquote>
</aside>
<p>If I understand what you want here correctly, you would first select out the <code>idx</code><br>
slice of dimension <code>d_n</code>.  Then <code>reshape()</code> dimensions <code>d_0, ..., d_{n - 2}</code><br>
into a single “batch” dimension, while using the <code>reshape()</code> to add a singleton<br>
channels dimension.  Apply the <code>Conv1d</code> and <code>reshape()</code> again to get your<br>
desired output.</p>
<p>Best.</p>
<p>K. Frank</p>

imperialn00b: <p>Thanks! This worked:</p>
<pre><code class="lang-auto">k = torch.nn.Conv1d(1, 1, kernel_size=N, bias=False)
with torch.no_grad():
    k.weight = torch.nn.Parameter(torch.ones_like(k.weight) * 1/N)
    output = k(input[..., idx].unsqueeze(1)).squeeze(1)
</code></pre>