title: Numpy sum elements in a multi-dimensional array according to indices
body: <p>I am dealing with a very large multi-dimensional data , but let me take a 2D array for example. Given a value array that is changing every iteration,</p>
  <pre><code>arr = np.array([[ 1, 2, 3, 4, 5], [5, 6, 7, 8, 9]]) # a*b
  </code></pre>
  <p>and an index array that is fixed all the time.</p>
  <pre><code>idx = np.array([[[0, 1, 1], [-1, -1, -1]],
                  [[5, 1, 3], [1, -1, -1]]]) # n*h*w, where n = a*b,
  </code></pre>
  <p>Here -1 means no index will be applied. And I wish to get a result</p>
  <pre><code>res = np.array([[1+2+2, 0],
                  [5+2+4, 2]]) # h*w
  </code></pre>
  <p>In real practice, I am doing with a very large 3D tensor (n ~ trillions), with a very sparse idx (i.e. lots of -1). As <code>idx</code> is fixed, my current solution is to pre-compute a n*(h*w) array <code>index_tensor</code> by filling 0 and 1, and then do</p>
  <pre><code>tmp = arr.reshape(1, n)
  res = (tmp @ index_tensor).reshape([h,w])
  </code></pre>
  <p>It works fine but takes a huge memory to store the <code>index_tensor</code>. Is there any approach that I can take the advantage of the sparsity and unchangeableness of <code>idx</code> to reduce the memory cost and keep a fair running speed in python (using numpy or pytorch would be the best)? Thanks in advance!</p>