torch.einsum 是一种强大的张量操作工具，可以通过爱因斯坦求和约定（Einstein summation convention）来简洁地表示复杂的张量运算。通过它，我们可以高效地计算矩阵乘法、转置、点积、外积等操作。

以下是关于如何使用 torch.einsum 计算两个四维张量在第三维度上的外积的解释，以及计算原理的详细说明。

---

计算两个四维张量在第三维度上的外积

假设有两个四维张量：

• A 的形状是 (N, M, P, Q)
• B 的形状是 (N, M, R, S)

我们希望在第三维度（P 和 R）上计算它们的外积，结果张量的形状应为 (N, M, P, R, Q, S)。

代码示例

可以通过以下方式实现：

import torch

# 假设 A 和 B 是随机初始化的张量
A = torch.randn(2, 3, 4, 5)  # 形状 (N=2, M=3, P=4, Q=5)
B = torch.randn(2, 3, 6, 7)  # 形状 (N=2, M=3, R=6, S=7)

# 使用 torch.einsum 计算外积
result = torch.einsum('nmij,nmok->nmijok', A, B)

print(result.shape)  # 输出 (2, 3, 4, 6, 5, 7)

详细计算原理

1. 爱因斯坦求和约定