torch.broadcast_tensors()原理解析

在a, b = torch.broadcast_tensors(x, y)中,是将x与y的形状“黏合”起来，从而组成两个形状相同的tensor：a、b，其中最主要要注意的点是，x的第一维（最后一维）必须与y的最后一维（第一维）相同且为1，例如：

x.shape:torch.Size([1, 2, 4]), 
y.shape:torch.Size([2, 2, 1])
a, b = torch.broadcast_tensors(x, y)

最终会得到的形状为

a.shape:torch.Size([2, 2, 4]), 
b.shape:torch.Size([2, 2, 4])

最终代码且运行结果为

import torch

x = torch.arange(8).view(1, 2, 4)
y = torch.arange(4).view(2, 2, 1)
a, b = torch.broadcast_tensors(y, x)
print('x:{},\ny:{}'.format(x, y))
print('x.shape:{}, \ny.shape:{}'.format(x.shape, y.shape))
print('a:{},\nb:{}'.format(a, b))
print('a.shape:{}, \nb.shape:{}'.format(a.shape, b.shape))

结果为

x:tensor([[[0, 1, 2, 3],
         [4, 5, 6, 7]]]),
y:tensor([[[0],
         [1]],

        [[2],
         [3]]])
x.shape:torch.Size([1, 2, 4]), 
y.shape:torch.Size([2, 2, 1])
a:tensor([[[0, 1, 2, 3],
         [4, 5, 6, 7]],

        [[0, 1, 2, 3],
         [4, 5, 6, 7]]]),
b:tensor([[[0, 0, 0, 0],
         [1, 1, 1, 1]],

        [[2, 2, 2, 2],
         [3, 3, 3, 3]]])
a.shape:torch.Size([2, 2, 4]), 
b.shape:torch.Size([2, 2, 4])