Fancy indexing in pytorch

I need to select a “view” of a Tensor, and assigning other values to it.
In NumPy I do something like this:

ar = np.random.rand(2, 3, 5)
xs = np.array([0, 1, 0, 0, 1])
ys = np.array([0, 1, 2, 3, 4])
res = np.random.rand(5,3)
ar[xs, :, ys] = res

The fancy indexing views the array ar, selecting a subset of size (5,3), then I insert into it the res values.
How can I do the same thing in pytorch. More in general how do I perform fancy indexing in Pytorch?

I have taken a look to the gather, select_index, and scatter_ functions, but none of them seems to work for this example.

3 Likes

Did you find an answer?

@AreTor’s snippet looks to be an example of NumPy “Advanced Indexing” with non-adjacent advanced indexing objects, as described here: https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#combining-advanced-and-basic-indexing. Unfortunately, we do not have advanced indexing in PyTorch for non-adjacent advanced indexing objects (i.e. in this case the fact that xs and ys are not adjacent.