title: PyTorch vectorise sum lookup quantity into buckets
body: <p>Using PyTorch, I have figured out the following code for calculating totals of an item's property by some "bucket index":</p>
  
  <pre><code>DATASET_SIZE = 10
  NUM_BUCKETS = 4
  bucket_assignment = torch.tensor([0,1,2,3,0,1,2,3,0,1], dtype = torch.long)
  values_to_add = torch.tensor([1,2,3,4,5,6,7,8,9,10], dtype = torch.float)
  buckets = torch.zeros(NUM_BUCKETS, dtype = torch.float)
  buckets.index_add_(0, bucket_assignment, values_to_add)
  
  # Buckets is now tensor([15., 18., 10., 12.])
  </code></pre>
  
  <p>In my case this is specifically to check allocation bounds on a problem, and later code checks that no bucket is under- or over- allocated.</p>
  
  <p>I would like to check multiple different possible assignments at once (and later pick a best option, code not shown). I thought I could do this by adding another dimension to <code>bucket_assignment</code> plus to <code>buckets</code> and have each row be a different set of assignments. However, this does not work as intended, because the second argument of <code>index_add_</code> must be a simple vector, I cannot pass in any higher rank tensor.</p>
  
  <p>E.g.</p>
  
  <pre><code>BATCH_SIZE = 2
  DATASET_SIZE = 5
  NUM_BUCKETS = 3
  bucket_assignment = torch.tensor([[0,1,2,0,1], [1,1,1,2,1]], dtype = torch.long)
  values_to_add = torch.tensor([1,2,3,4,5], dtype = torch.float)
  buckets = torch.zeros(BATCH_SIZE, NUM_BUCKETS, dtype = torch.float)
  buckets.index_add_(0, bucket_assignment, values_to_add)
  </code></pre>
  
  <p>I would like to get this result:</p>
  
  <pre><code>tensor([[5., 7., 3.], [ 0., 11.,  4.]])
  </code></pre>
  
  <p>Instead, I get an error:</p>
  
  <pre><code>RuntimeError: invalid argument 3: Index is supposed to be a vector at ../aten/src/TH/generic/THTensorEvenMoreMath.cpp:733
  </code></pre>
  
  <p>That's not unexpected due to the limitations of <code>.index_add</code>, but I don't know how to progress.</p>
  
  <p>I am not sure what other approach would allow me to solve this problem in PyTorch - is there some other <code>torch</code> method I could use that would allow me to achieve the same thing. The main goal here is vectorisation and avoiding loops in Python, as in reality the batch sizes are large and I will be taking advantage of GPU acceleration.</p>