I would like to do a += operation at specified columns of a 2d matrix, where the column indices are in another 2D matrix.

Specifically, is there a more efficient way to do the operation below.

import numpy as np

N, V, D, B = 5, 25000, 300, 512

grad = np.zeros((D, V))

# col_idx columns and rows can have duplicate indices, so 
# as += applied at the vector level leads to only a single 
# instance of the index working.
col_idx = np.random.randint(0, V, size=(N, B))

# force at least one example of duplicates within a column.
col_idx[:, 0] = np.array([0, 100, V, 0, V])

values = np.random.normal(size=(D, B))

for b in range(B):
  for row in range(col_idx.shape[0]): 
    grad[:, col_idx[row, b]] += values[:, b]

np.add.at is supposed to be useful when += doesn’t work, but couldn’t quite figure out how to use it in this case.