torch.einsum()

`torch.einsum()` 是 PyTorch 中的一个函数，它允许对张量执行元素级别的矩阵运算，支持定义自定义的缩写字符串来表示张量的维度并指定操作方式。这个函数非常适合进行高效的多维数组操作，特别是对于那些具有特定模式的求和、卷积和相似操作。

例如，在给定的代码示例中[^1]:

# 定义张量 A 和 x
A = torch.randn(3, 4)  # 形状 (V, N)
x = torch.randn(2, 5, 4, 6)  # 形状 (B, F, N, T)

# 使用 einsum 求 A 对 x 的特定维度的乘法和拼接
result = torch.einsum('vn,bfnt->bfvt', (A, x))  # 'vn' 表示 A 的第一个维度与 x 的第二个维度相乘
# 'bfn' 和 't' 合并成新的 'bfvt'，意味着将 x 的最后一个维度和 A 的第二个维度的结果连接起来

# 打印结果形状
print(result.shape)  # 输出 (2, 5, 3, 6)

`einsum` 可以用于各种复杂的张量计算，如矩阵乘法（`'ij,jk->ik'`）、向量化操作（`'n->()'`），甚至于自定义的数学表达式。其灵活性使得在深度学习和科学计算中非常有用。