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

```
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:
print('didnt work with' , mdp_data('sa_s'))
```

Given that `vp_sa_s`

is size `(10,5,5)`

and each value is a valid index vp i.e 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!