Снижение тензора на основе индекса вектора

Question

В качестве примера у меня есть 2 тензора: A = [1;2;3;4;5;6;7] и B = [2;3;2]. Идея состоит в том, что я хочу уменьшить A на основе B - так, чтобы значения B представляли, как суммировать значения A - так, чтобы B = [2;3;2] означало, что уменьшенное A будет суммой первых 2 значений, следующих 3 и последних 2: A' = [(1+2);(3+4+5);(6+7)]. Очевидно, что сумма B всегда должна быть равна длине A. Я пытаюсь сделать это настолько эффективно, насколько это возможно - желательно, чтобы определенные функции или матричные операции содержались в pytorch / python. Спасибо!

Answer 1

Вот решение.

• Сначала мы создаем массив индексов B_idx с таким же размером A.
• Затем накапливать (добавлять) все элементы в A на основе индексов B_idx, используя index_add_.

A = torch.arange(1, 8) 
B = torch.tensor([2, 3, 2])

B_idx = [idx.repeat(times) for idx, times in zip(torch.arange(len(B)), B)]
B_idx = torch.cat(B_idx) # tensor([0, 0, 1, 1, 1, 2, 2])

A_sum = torch.zeros_like(B)
A_sum.index_add_(dim=0, index=B_idx, source=A)
print(A_sum) # tensor([ 3, 12, 13])