So i have a tensor with a certain amount of dimensions A filled with values and a Tensor C with the same dimensions (but without the last) with indices for A’s last dimension. I want to pick the values of A based on C such that a Tensor with the shape of C is returned.
Example:
A has shape (7, 2, 4), C has shape (7, 2).
A:
[[[-10.0491, -4.9780, -3.0346, -1.0746],
[ 1.1812, 1.8627, 3.2540, 5.5354]],
[[ -9.9464, -4.9588, -2.9927, -0.8602],
[ 1.0148, 1.9898, 2.7656, 4.7714]],
[[ -9.7778, -5.0038, -3.0378, -1.0750],
[ 1.0365, 2.1866, 2.8971, 4.8669]],
[[-10.1701, -4.8115, -3.0066, -1.0485],
[ 0.7645, 1.8798, 3.0735, 4.7153]],
[[ -9.8422, -4.9419, -3.1802, -0.8486],
[ 0.8864, 1.7320, 3.3117, 4.8196]],
[[-10.0794, -4.9817, -2.9976, -0.9717],
[ 0.6711, 1.9173, 2.8586, 4.9084]],
[[-10.0262, -5.1335, -2.9970, -0.9397],
[ 1.2652, 1.9704, 3.0415, 5.8505]]]
I:
[[0, 1],
[0, 3],
[0, 1],
[0, 3],
[1, 0],
[2, 1],
[2, 1]]
The wanted resulting tensor B with shape (7, 2) is:
[[-10.0491, 1.8627],
[ -9.9464, 4.7714],
[ -9.7778, 2.1866],
[-10.1701, 4.7153],
[ -4.9419, 0.8864],
[ -2.9976, 1.9173],
[ -2.9970, 1.9704]]
I want this to be possible for A being n-dimensional and C being (n-1)-dimensional.
I had this solution for n=3:
b = a[torch.arange(c.shape[0]).unsqueeze(1), torch.arange(c.shape[1]), c]
This will obviously cause problems for a differnt amount of dimensions.
Thank you for help!