How to batch indexing tensors in Julia


#1

e.g A = rand(30, 2, 30) and now I have a batch of index [1, 2, 1, 1, 1, 2], I want to have the result of this batched indexing res of size (30, 6, 30), where

res[:, 1, :] == A[:, 1, :]

res[:, 2, :] == A[:, 2, :]
res[:, 3, :] == A[:, 1, :]

Is there any builtin function that can do this?


#2

I think A[:, [1, 2, 1, 1, 1, 2], :] could be what you are looking for.