本文旨在深入解析 NumPy 中 `einsum` 函数的用法，通过具体示例和代码演示，帮助读者理解其在张量运算中的作用，并掌握利用 `einsum` 实现高效、灵活的张量操作的方法。文章将从基本概念入手，逐步剖析 `einsum` 的运算规则，并提供等效的循环实现，以便读者更好地理解其内部机制。
np.einsum 是 NumPy 中一个强大的函数，用于执行爱因斯坦求和约定。它允许用户以简洁而灵活的方式表达各种张量运算，例如矩阵乘法、张量缩并、求迹等。理解 einsum 的工作原理对于高效处理多维数组至关重要。
einsum 的基本语法
einsum 的基本语法如下：
np.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False)
其中，subscripts 是一个字符串，用于指定张量的维度以及如何进行求和。operands 则是要进行运算的张量。
理解 subscripts 字符串
subscripts 字符串的核心在于定义了输入张量的维度标签，以及输出张量的维度标签。相同的维度标签表示在该维度上进行求和（缩并）。
例如，'ijk,jil->kl' 表示：
- 第一个输入张量 a 的维度是 ijk。
- 第二个输入张量 b 的维度是 jil。
- 输出张量的维度是 kl。
- 由于 i 和 j 在输入张量中出现，但在输出张量中没有出现，因此在 i 和 j 维度上进行求和。
示例分析
考虑以下示例：
import numpy as np
a = np.arange(8.).reshape(4, 2, 1)
b = np.arange(16.).reshape(2, 4, 2)
result = np.einsum('ijk,jil->kl', a, b)
print(result)为了更好地理解 einsum 的运算过程，我们可以将其分解为更细粒度的步骤。首先，考虑不进行任何求和的情况：
result_no_sum = np.einsum('ijk,jil->ijkl', a, b)
print(result_no_sum)这个操作将生成一个形状为 (4, 2, 1, 2) 的张量，其中每个元素是 a 和 b 中对应元素的乘积，而没有进行任何求和。
接下来，为了得到原始 einsum 的结果，我们需要在轴 0 和轴 1 上进行求和：
result_sum_1 = result_no_sum.sum(axis=1) print(result_sum_1) result_sum_0 = result_sum_1.sum(axis=0) print(result_sum_0)
result_sum_0 的结果与 np.einsum('ijk,jil->kl', a, b) 的结果相同。
等效的循环实现
为了更深入地理解 einsum 的运算过程，我们可以使用循环来实现相同的功能：
def sum_array(A, B):
i_len, j_len, k_len = A.shape
_, _, l_len = B.shape
ret = np.zeros((k_len, l_len))
for i in range(i_len):
for j in range(j_len):
for k in range(k_len):
for l in range(l_len):
ret[k, l] += A[i, j, k] * B[j, i, l]
return ret
result_loop = sum_array(a, b)
print(result_loop)这段代码的逻辑与 einsum('ijk,jil->kl', a, b) 完全相同。通过显式循环，我们可以清楚地看到每个元素的乘积以及求和的过程。
注意事项和总结
- einsum 的性能通常优于显式循环，尤其是在处理大型张量时。
- einsum 的灵活性使其能够表达各种复杂的张量运算，但同时也需要仔细理解 subscripts 字符串的含义。
- 通过将 einsum 分解为更细粒度的步骤，例如先不进行求和，然后再进行求和，可以帮助理解其运算过程。
- 使用等效的循环实现可以更深入地理解 einsum 的内部机制。
总而言之，np.einsum 是一个强大的工具，可以帮助我们高效地处理张量运算。通过理解其基本语法和运算规则，我们可以充分利用其灵活性和性能优势。
