PyTorch: 2D-тензор индекса с 2D-тензором индексов строк

Question

У меня есть тензор факела a формы (x, n) и еще один тензор b формы (y, n), где y <= x.каждый столбец b содержит последовательность индексов строк для a, и я хотел бы как-то индексировать a с помощью b так, чтобы я получил тензор формы (y, n), в которомi-й столбец содержит a[:, i][b[:, i]] (не совсем уверен, правильно ли это выражено).

Вот пример (где x = 5, y = 3 и n = 4):

import torch

a = torch.Tensor(
    [[0.1, 0.2, 0.3, 0.4],
     [0.6, 0.7, 0.8, 0.9],
     [1.1, 1.2, 1.3, 1.4],
     [1.6, 1.7, 1.8, 1.9],
     [2.1, 2.2, 2.3, 2.4]]
)

b = torch.LongTensor(
    [[0, 3, 1, 2],
     [2, 2, 2, 0],
     [1, 1, 0, 4]]
)

# How do I get from a and b to c
# (so that I can also assign to those elements in a)?

c = torch.Tensor(
    [[0.1, 1.7, 0.8, 1.4],
     [1.1, 1.2, 1.3, 0.4],
     [0.6, 0.7, 0.3, 2.4]]
)

Я не могу разобраться с этим.Мне нужен метод, который не даст тензор c, но также позволит мне назначить тензор той же формы, что и c, элементам a, из которых состоит c.

Answer 1

Я пытаюсь использовать index_select, но он поддерживает только массив 1-dim для индекса.

bt = b.transpose(0, 1)
at = a.transpose(0, 1)
ct = [torch.index_select(at[i], dim=0, index=bt[i]) for i in range(len(at))]
c  = torch.stack(ct).transpose(0, 1)
print(c)
"""
tensor([[0.1000, 1.7000, 0.8000, 1.4000],
        [1.1000, 1.2000, 1.3000, 0.4000],
        [0.6000, 0.7000, 0.3000, 2.4000]])
"""

Возможно, это не лучшее решение, но надеюсь, что это поможет вам, по крайней мере.