How to speed up slicing in python, not using the for loop
I'm trying to speed up the following python code:

import torch
import numpy as np

A = torch.zeros(11, 16, 64)
B = torch.randn(11, 9, 64)

indices = np.random.randint(0,9,(11,16))

for i in range(len(A)):
    A[i,:,:] = B[i,indices[i],:]


Is there a nice way not to use the for loop? This way, it is really slow, especially when dealing with the big data . The indices is the pre-defined 2-dim matrix with size (11,16). What I need is to assign the elements of B to A according to the order of indices. After the speeding up, the result of A should be exactly the same with mine resulted A. Thanks!

You can use multiple mult-dimensional indices but they need to be the same size or broadcastable. So for example

# create a (11, 1) range array that broadcasts with indices which is (11, 16)
indices0 = np.expand_dims(np.arange(indices.shape[0]), 1)
A = B[indices0, indices, :]


Since broadcasting can be confusing I'll try to explain this a little. Basically you want indices0 and indices to be the same size and represent pairs of indices of B. The first index will be stored in indices0 and the second will be stored in indices in corresponding locations. Broadcasting implicitly repeats the columns of indices0 to make it the same shape as indices and can often be faster than constructing the full sized indices0.

In case it helps here are some more verbose examples demonstrating why this works:

import torch
import numpy as np

B = torch.randn(11, 9, 64)
indices = np.random.randint(0,9,(11,16))

# constructing indices0 more verbosely (and slower) for demonstration purposes
a0, a1 = indices.shape
a2 = B.shape[2]

# construct a complete indices0 the slow way, the same size as indices
indices0 = np.empty((a0, a1), dtype=np.int32)
for i in range(a0):
    for j in range(a1):
        indices0[i,j] = i

# version 1 (nothing complicated happening here but very slow)
A1 = torch.empty(a0, a1, a2, dtype=B.dtype)
for i in range(a0):
    for j in range(a1):
        A1[i,j,:] = B[indices0[i,j], indices[i,j], :]

# version 2 (using advanced indexing without broadcasting)
A2 = B[indices0, indices, :]

# version 3 (with broadcasting)
# remove repeated columns leaving indices0 as (11, 1) the same state as above
indices0 = indices0[:, :1]
# broadcasting implicitly repeats columns of indices0 to match indices
A3 = B[indices0, indices, :]

# version 4 (your method)
A4 = torch.empty(a0, a1, a2, dtype=B.dtype)
for i in range(a0):
    A4[i,:,:] = B[i,indices[i],:]

# compare everything    
error = torch.sum(torch.abs(A1 - A2)).item() + \
        torch.sum(torch.abs(A2 - A3)).item() + \
        torch.sum(torch.abs(A3 - A4)).item()
print('Error:', error)


which prints

Error: 0.0


demonstrating that all these methods are equivalent.



Also, if you wanted to stay in the PyTorch framework and indices were a torch.LongTensor instead of a numpy.ndarray then you could use

indices0 = torch.arange(indices.shape[0]).unsqueeze(1)
A = B[indices0, indices, :]

