PyTorch: пакетное внешнее дополнение

Question

У меня есть два тензора PyTorch: A и B, оба имеют форму (b, c, 3). Я хочу сделать внешний продукт C из A и B так, чтобы полученная фигура была (b, c, 3, 3), и заменить операцию умножения на сложение. Как мне это сделать?

Answer 1

Вы можете добавить соответствующее одноэлементное измерение:

C = A[..., None] + B[..., None, :]

Например, с размерами партии и канала, равными 1 (b=1, c=1):

import torch
A = torch.tensor([[[1, 2, 3.]]]) 
B = torch.tensor([[[4., 5., 6.]]])
A[..., None] + B[..., None, :]
Out[ ]:
tensor([[[[5., 6., 7.],
          [6., 7., 8.],
          [7., 8., 9.]]]])