einsum（爱因斯坦求和）

在Python中，特别是使用NumPy或PyTorch等科学计算库时，爱因斯坦求和约定（Einstein summation convention）是一种强大的工具，它允许你以简洁的方式指定多维数组（张量）之间的操作，如求和、乘法和广播。虽然Python本身并不直接支持爱因斯坦求和约定的语法，但一些库（如NumPy的einsum函数和PyTorch的einsum操作）提供了对它的支持。

NumPy中的einsum

NumPy的einsum函数允许你指定数组操作的计算方式，通过一种紧凑的字符串表示法。这个字符串描述了输入数组的形状、操作的维度以及如何对它们进行求和或广播。

基本语法

numpy.einsum(subscripts, *operands, **kwargs)

• subscripts：一个字符串，描述了操作。它指定了输入数组的每个维度以及如何在输出数组中组合它们。
• operands：一个或多个数组，是操作的输入。
• **kwargs：其他关键字参数，如out（指定输出数组），dtype（指定输出数组的数据类型），order（指定数组的内存布局）等。

示例

1. 矩阵乘法

import numpy as np

a = np.random.rand(3, 4)
b = np.random.rand(4, 5)

# 使用einsum进行矩阵乘法
c = np.einsum('ij,jk->ik', a, b)

# 等同于使用numpy的dot函数
# c = np.dot(a, b)

在这个例子中，'ij,jk->ik'指定了操作：i和j是a的维度，j和k是b的维度，输出c的维度是i和k，其中j维度被求和（这是爱因斯坦求和约定的核心）。

2. 计算内积

a = np.random.rand(3)
b = np.random.rand(3)

# 使用einsum计算内积
dot_product = np.einsum('i,i->', a, b)

# 等同于使用numpy的dot函数（对于一维数组）
# dot_product = np.dot(a, b)

在这个例子中，'i,i->'指定了操作：i是a和b的公共维度，输出是一个标量（没有维度），因为所有维度都被求和了。

PyTorch中的einsum

PyTorch也提供了一个einsum操作，其语法和功能与NumPy的einsum非常相似。

示例

import torch

a = torch.randn(3, 4)
b = torch.randn(4, 5)

# 使用einsum进行矩阵乘法
c = torch.einsum('ij,jk->ik', [a, b])

# 等同于使用torch的matmul函数
# c = torch.matmul(a, b)

在这个PyTorch的例子中，einsum的调用方式与NumPy非常相似，但是注意输入数组是作为列表传递给einsum的。

总结

爱因斯坦求和约定是一种强大的工具，它允许你以一种非常紧凑和灵活的方式指定多维数组之间的操作。虽然Python本身不直接支持这种语法，但NumPy和PyTorch等库提供了对它的支持，使得你可以利用这种约定来编写更简洁、更高效的代码。