Numpy / PyTorch - how to assign value with indices in different dimensions?

Question

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]])

Answer 1

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

Answer 2

Since this question is tagged with PyTorch, here is a PyTorch equivalent solution just for the sake of completeness.

# make a copy of the inputs from numpy array
In [188]: at = torch.tensor(a)
In [189]: at_idxs = torch.tensor(a_indices)

We will make use of tensor.scatter_(...) to do the replacement in-place. So, let's prepare the inputs first.

torch.scatter() API expects the replacement value (here 100) to be a tensor and of the same shape as the indices tensor. Thus, we have to create a tensor filled with a value of 100 and of the shape (2, 2) since the indices tensor at_idxs is of that shape. So,

In [190]: replace_val = 100 * torch.ones(at_idxs.shape, dtype=torch.long)    
In [191]: replace_val 
Out[191]: 
tensor([[100, 100],
        [100, 100]])

Perform the in-place replacement now:

# fill the values along dimension 1
In [198]: at.scatter_(1, at_idxs, replace_val)
Out[198]: 
tensor([[100,   2, 100],
        [  4, 100, 100]])