Is there a Numpy or pyTorch function for this code?

Question

Basically is there a Numpy or PyTorch function that does this:

vp_sa_s=mdp_data['sa_s'].detach().clone()
dims = vp_sa_s.size()
for i in range(dims[0]):
    for j in range(dims[1]):
        for k in range(dims[2]):
            # to mimic matlab functionality: vp(mdp_data.sa_s)
            try:
                vp_sa_s[i,j,k] = vp[mdp_data['sa_s'][i,j,k]]
            except:
                pass

Given that vp_sa_s is size (10,5,5) and each value is a valid index vp ie in range 0-9. vp is size (10,1) with a bunch of random values.

Matlab do it elegantly and quickly with vp(mdp_data.sa_s) which will form a new (10,5,5) matrix. If all values in mdp_data.sa_s are 1, the result would be a (10,5,5) tensor with each value being the 1st value in vp .

Does a function or method that exists that can achieve this in less than O(N^3) time as the above code is terribly inefficient.

Thanks!

Answer 1

What is wrong with

result = vp[vp_sa_s, 0]

note that since your vp is of shape (10, 1) (it has a trailing singleton dimension) you need to add the , 0] index in the assignment to get rid of this extra dimension.