pytorch 矩阵运算实现“笛卡尔加”

问题定义见下图：给定A，B，求A“笛卡尔加”B（笛卡尔加类似于笛卡尔积，只是在合并时做了加法运算）

用for语句解决这个问题当然非常简单，但是for的性能损失太大（我在神经网络里，分别用for和矩阵运算实现了“笛卡尔加”，最终for的时间开销为8h+，矩阵运算为5min+），记录一下怎么用矩阵运算解决：

step

扩充A，B为（3,3,4）维度，得到_A , _B.
_A=([A1,A2,A3],[A1,A2,A3],[A1,A2,A3]) _B同理。
将_A的前两维转置后+_B即可得到AB的“笛卡尔加”

code

以此为例

		self.W_2 = torch.nn.Parameter(torch.Tensor(in_channels, out_channels))
        self.U_2 = torch.nn.Parameter(torch