I have a 5 dimensional tensor S of shape (bs,a,b,c,d) which can be seen as a batch of (a x b) 2D matrices of (c x d) 2D matrices. Let say bs =1 and a = b = c = d = 3 for simplicity.
I have a shape (b) vector (let say ind = [2,0,1]) that contains indices to columns of the inner c x d matrices. For each row i and column j of the outer a x b matrix, I need to gather the column ind[j] of the c x d matrix in an (a x b x c) tensor.
Example with S
tensor([[[[[ 0, 0, 0],
[ 0, 0, 0],
[ 0, 0, 0]],
[[1211, 1212, 1213],
[1221, 1222, 1223],
[1231, 1232, 1233]],
[[1311, 1312, 1313],
[1321, 1322, 1323],
[1331, 1332, 1333]]],
[[[1211, 1221, 1231],
[1212, 1222, 1232],
[1213, 1223, 1233]],
[[ 0, 0, 0],
[ 0, 0, 0],
[ 0, 0, 0]],
[[2311, 2312, 2313],
[2321, 2322, 2323],
[2331, 2332, 2333]]],
[[[1311, 1321, 1331],
[1312, 1322, 1332],
[1313, 1323, 1333]],
[[2311, 2321, 2331],
[2312, 2322, 2332],
[2313, 2323, 2333]],
[[ 0, 0, 0],
[ 0, 0, 0],
[ 0, 0, 0]]]]])
I found that
S[:,torch.arange(3)[None,:],torch.arange(3)[:,None],ind[None,:]]
does the job for this specific case
tensor([[[[ 0, 0, 0],
[1211, 1221, 1231],
[1312, 1322, 1332]],
[[1231, 1232, 1233],
[ 0, 0, 0],
[2312, 2322, 2332]],
[[1331, 1332, 1333],
[2311, 2312, 2313],
[ 0, 0, 0]]]])
does the job but I don’t get why and how broadcasting works here. My indices have all shape (1,3) or (3,1) and should broadcast to a (3,3) shape but the result has shape (3,3,3) as required. I don’t get why.
What would be a general solution for arbitrary bs,a,b,c,d?
T.