Sum over various subsets of a tensor

James_Hickman: <p>This may have already been addressed, but I did some googling and couldn’t find a solution. I’d like to compute various sums from unequal sized subsets of a given tensor (or more precisely from a column vector) where the summing index boundaries are defined by a list (or tensor) and then to have the operation return a tensor of these sums  (without using a for loop)</p>
<p>torch.split does almost exactly what I want but it returns a list, for example</p>
<p>a=torch.tensor([[1,2,3,4,5,6,7,8]])</p>
<p>print(torch.split(a,[2,2,4],dim=1))<br>
returns the following lists<br>
(tensor([[1, 2]]), tensor([[3, 4]]), tensor([[5, 6, 7, 8]]))</p>
<p>what I want is a tensor whose elements are the sums of those individual lists<br>
i.e tensor([[3,7,26]]) .<br>
Is it possible to do this in a vectorized way using torch operations? also I’d like to do the same thing but with a product over the subsets and return a tensor of the products.</p>
<p>Much thanks in advance!</p>

James_Hickman: <p>I’m still looking for a better solution but if someone is having a similar problem, I sort of figured out how to do it with the help of ideas in the following link:</p><aside class="quote" data-post="3" data-topic="7184">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/ac91a4/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/efficient-way-of-calculating-sum-of-unequal-sized-chunks-of-tensor/7184/3">Efficient way of calculating sum of unequal sized chunks of tensor</a> 
  </div>
  <blockquote>
    Thanks for your reply. I am actually doing this on GPU. I think the only way is to pad it and do sum along axis. I would assume if I am doing for loop anyway, only do padding with in the for loop would be faster? In any case I will try it out. Thanks a lot!
  </blockquote>
</aside>

<p>This isn’t the most efficient method because a lot of the terms in cumsum (and cumprod) are unnecessary, if someone finds a more efficient way to do this then please post here but the following seems to work at least for the summation: (NOTE: this has problems with cumprod, for example if  the array is long then the cumulative product can blow up to INF or if there are any zeros in the array, then all subsequent products are zero)</p>
<p>PYTORCH SOLUTION</p>
<p>FOR THE SUM<br>
a=torch.tensor([[1,2,3,4,5,6,7,8]])<br>
b=a.cumsum(1) <span class="hashtag-raw">#cumulative</span> sum over row<br>
c=b.gather(1, torch.tensor([[1,3,7]])) <span class="hashtag-raw">#select</span> relevant terms<br>
d=torch.cat( (torch.tensor([[0]]), b.gather(1, torch.tensor([[1,3]]))),1) <span class="hashtag-raw">#select</span> relevant terms<br>
print(c,d,c-d)</p>
<p>returns</p>
<p>tensor([[ 3, 10, 36]]) tensor([[ 0,  3, 10]]) tensor([[ 3,  7, 26]])</p>
<p>FOR THE PRODUCT<br>
a=torch.tensor([[1,2,3,4,5,6,7,8]])<br>
b=a.cumprod(1) <span class="hashtag-raw">#cumulative</span> sum over row<br>
c=b.gather(1, torch.tensor([[1,3,7]])) <span class="hashtag-raw">#select</span> relevant terms<br>
d=torch.cat( (torch.tensor([[1]]), b.gather(1, torch.tensor([[1,3]]))),1) <span class="hashtag-raw">#select</span> relevant terms<br>
print(c,d,c/d)</p>
<p>returns<br>
tensor([[    2,    24, 40320]]) tensor([[ 1,  2, 24]]) tensor([[   2,   12, 1680]])</p>
<p>NUMPY SOLUTION<br>
a = np.array([[1,2,3,4,5,6,7,8]])<br>
b=np.add.reduceat(a, [0,2,4], axis=1)<br>
print(b,type(b),b.shape)</p>
<p>returns</p>
<p>[[ 3  7 26]] &lt;class ‘numpy.ndarray’&gt; (1, 3)</p>

Remorax: <p>Thanks for the answer <a class="mention" href="/u/james_hickman">@James_Hickman</a>!</p>
<p>I am also facing this problem right now, and am looking for a more efficient solution since I am dealing with large tensors.</p>
<p>Bumping this post in case anyone knows of a better solution now, because of the improvements in Pytorch over the last two years.</p>

Bubas3000: <p><a class="mention" href="/u/remorax">@Remorax</a> Did you find a more efficient solution? I have the same problem.</p>
<p>Thank you.</p>

prcosta: <p>Perhaps this is helpful.</p>
<p>Based on the solution of <a class="mention" href="/u/james_hickman">@James_Hickman</a> :</p>
<pre><code class="lang-python">a = torch.tensor([[1,2,3,4,5,6,7,8]])
b = a.cumsum(1)
c = b.gather(1, torch.tensor([[1,3,7]])) #select relevant terms
c = torch.cat([torch.zeros(1, 1), c], dim=-1) #start the sum with zeros
res = c.unsqueeze(2) - c.unsqueeze(1)
torch.diagonal(res, offset=-1, dim1=1, dim2=2)
</code></pre>
<p>returns</p>
<pre><code class="lang-python">tensor([[ 3.,  7., 26.]])
</code></pre>

DillonZChen: <p>The scatter_add function in the pytorch-scatter package seems to do exactly what you want. See <a href="https://pytorch-scatter.readthedocs.io/en/1.3.0/functions/add.html" class="inline-onebox" rel="noopener nofollow ugc">Scatter Add — pytorch_scatter 1.3.0 documentation</a></p>

hshuai97: <p>this is helpful to me. thank you very much!</p>

aRI0U: <p>Finally found a nice solution to solve this problem, which consists in basically multiplying the input tensor by a one-hot corresponding to the indices of the elements.</p>
<p>Assume you have a tensor <code>input</code> of size <code>Nxd</code> (i.e. N <code>d</code>-dimensional elements) and a tensor <code>encoding_indices</code> of size <code>N</code> s.t. <code>encoding_indices[i]</code> is the id of the subset in which <code>input[i]</code> is. Let <code>num_codes</code> be the number of subsets.</p>
<p>Then, the following script computes the sum of all vectors aggregated per subset:</p>
<pre><code class="lang-python">one_hot = F.one_hot(encoding_indices, num_classes=num_codes).float()
sum_per_subset = torch.mm(one_hot.T, inputs)
</code></pre>
<p>Since this script implies a matrix product of size <code>nxN</code> times <code>Nxd</code> implying a very sparse matrix, this script may be optimized by using a sparse tensor to encode the one-hot vector:</p>
<pre><code class="lang-python">num_inputs = inputs.size(0)
indices = torch.stack((
    encoding_indices,
    torch.arange(num_inputs, device=encoding_indices.device)
))
values = torch.ones_like(encoding_indices, dtype=torch.float)
one_hot = torch.sparse_coo_tensor(indices, values, size=(num_codes, num_inputs))
sums_per_subset = torch.mm(one_hot, inputs)
</code></pre>
<p>Here is a performance comparison on CPU and GPU:</p>
<pre><code class="lang-python">cpu
f1: 1.087107 s  # naive for loop
f2: 0.087390 s  # dense matrix multiplication
f3: 0.016435 s  # sparse matrix multiplication

cuda:0
f1: 0.117819 s
f2: 0.010617 s
f3: 0.002403 s
</code></pre>
<p>Results obtained running the following script:</p>
<pre><code class="lang-python">import timeit

import torch
import torch.nn.functional as F


def f1(inputs, encoding_indices, num_codes):
    sum_inputs = torch.zeros(num_codes, inputs.size(1), device=inputs.device)
    for i in range(num_codes):
        sum_inputs[i] = inputs[encoding_indices == i].sum(dim=0)
    return sum_inputs


def f2(inputs, encoding_indices, num_codes):
    one_hot = F.one_hot(encoding_indices, num_classes=num_codes).float()
    return torch.mm(one_hot.T, inputs)


def f3(inputs, encoding_indices, num_codes):
    num_inputs = inputs.size(0)
    indices = torch.stack((
        encoding_indices,
        torch.arange(num_inputs, device=encoding_indices.device)
    ))
    values = torch.ones_like(encoding_indices, dtype=torch.float)
    one_hot = torch.sparse_coo_tensor(indices, values, size=(num_codes, num_inputs))
    return torch.mm(one_hot, inputs)


num_inputs = 32456
code_dim = 67
num_codes = 511

inputs = torch.rand(num_inputs, code_dim)
encoding_indices = torch.randint(0, num_codes, (num_inputs,))

for d in ["cpu", "cuda:0"]:
    device = torch.device(d)
    print(d)

    inputs = inputs.to(device)
    encoding_indices = encoding_indices.to(device)

    a1 = f1(inputs, encoding_indices, num_codes)
    a2 = f3(inputs, encoding_indices, num_codes)

    assert torch.allclose(a1, a2)

    for i in range(1, 4):
        print("f%d: %.6f s" % (
            i,
            timeit.timeit(
                f"f{i:d}(inputs, encoding_indices, num_codes)",
                number=100,
                globals=globals()
            ) / 100
        ))
</code></pre>
<p>I especially use this method for my implementation of EMA in VQ-VAE, since centroids are updated following the average of the inputs which are closest</p>