Объединение двух тензоров с разными размерами в Pytorch

Question

Возможно ли объединить два тензора с разными размерами без использования цикла for.

например, Tensor 1 имеет размеры (15, 200, 2048), а Tensor 2 имеет размеры (1, 200, 2048).Можно ли объединить 2-й тензор с 1-м тензором по всем 15 индексам 1-го измерения в 1-м тензоре (Трансляция 2-го тензора по 1-му измерению тензора 1 при одновременной конкатенации по 3-му измерению 1-го тензора)?Полученный тензор должен иметь размеры (15, 200, 4096).

Возможно ли это сделать без цикла for?

Answer 1

Вы можете выполнить вещание вручную (используя Tensor.expand()) до объединения (используя torch.cat()):

import torch

a = torch.randn(15, 200, 2048)
b = torch.randn(1, 200, 2048)

repeat_vals = repeat_vals = [a.shape[0] // b.shape[0]] + [-1] * (len(b.shape) - 1)
# or directly repeat_vals = (15, -1, -1) or (15, 200, 2048) if shapes are known and fixed...
res = torch.cat((a, b.expand(*repeat_vals)), dim=-1)
print(res.shape)
# torch.Size([15, 200, 4096])