Numpy / PyTorch - how to assign value with indices in different dimensions?
Suppose I have a matrix and some indices

a = np.array([[1, 2, 3], [4, 5, 6]])
a_indices = np.array([[0,2], [1,2]])


Is there any efficient way to achieve following operation?

for i in range(2):
    a[i, a_indices[i]] = 100

# a: np.array([[100, 2, 100], [4, 100, 100]])


Use np.put_along_axis -

In [111]: np.put_along_axis(a,a_indices,100,axis=1)

In [112]: a
Out[112]: 
array([[100,   2, 100],
       [  4, 100, 100]])


Alternaytively, if you want to do with the explicit way, i.e. integer-based indexing -

In [115]: a[np.arange(len(a_indices))[:,None], a_indices] = 100

