PyTorch中如何避免张量加减后压缩时的高额广播开销？

我有两个2D张量，它们会通过广播机制对齐维度。如果直接对它们做加减运算，会生成一个体积巨大的3D中间张量，但实际上我完全不需要这个中间张量——因为后续要立即对某一维度执行mean运算。

以下是用unsqueeze操作展示维度对齐逻辑的示例（实际中的张量都是2D的）：

x = torch.tensor(...)              # 形状: (batch , 1,  B)
y = torch.tensor(...)              # 形状: (1,    , A,  B)
out = torch.cos(x - y).mean(dim=2) # 最终输出形状: (batch, B)

目前考虑的解决方案方向：

• 代数简化：暂时还没找到可行的简化方式
• 借助PyTorch原生算子：这属于类余弦相似度计算，但和torch.cosine_similarity逻辑不同，我是将其应用在复数的.angle()运算场景中
• 编写高效循环的自定义C/CPython代码
• 其他可行方案？

内容的提问来源于stack exchange，提问作者Josh.F