`np.einsum` 提供了一种强大且简洁的方式来执行张量乘法和求和。本文旨在深入剖析 `einsum` 在处理两个张量（例如 `'ijk,jil->kl'`）时，其内部元素是如何进行乘法和求和的。我们将通过分解中间步骤和构建等效的显式循环两种方法，详细揭示 `einsum` 基于索引符号实现复杂张量运算的内在机制，帮助读者全面掌握其工作原理。
NumPy 中的 einsum 函数以其简洁而强大的张量操作能力而闻名，它允许用户通过爱因斯坦求和约定来表达复杂的张量运算，包括点积、外积、转置、张量收缩等。然而，对于初学者而言，理解 einsum 在执行如 np.einsum('ijk,jil->kl', a, b) 这类操作时，其内部元素是如何进行组合和求和的，可能是一个挑战。本文将通过具体的示例和两种不同的解析方法，详细阐述这一过程。
假设我们有两个 NumPy 张量 a 和 b：
import numpy as np
a = np.arange(8.).reshape(4,2,1)
b = np.arange(16.).reshape(2,4,2)
print("张量 a 的形状:", a.shape) # (4, 2, 1)
print("张量 b 的形状:", b.shape) # (2, 4, 2)我们的目标是理解 np.einsum('ijk,jil->kl', a, b) 的执行细节。这里的索引字符串 'ijk,jil->kl' 指示了以下操作：
这意味着：
为了理解元素是如何组合和求和的，我们可以首先修改 einsum 的输出索引，使其包含所有输入索引，从而避免任何隐式求和。这样，我们可以看到所有可能的元素乘积，然后再手动执行求和。
将 'ijk,jil->kl' 改为 'ijk,jil->ijkl'，这意味着输出张量将包含 i, j, k, l 四个维度，并且不会对任何索引进行求和。
# 查看所有元素乘积，不进行求和
intermediate_products = np.einsum('ijk,jil->ijkl', a, b)
print("中间乘积张量的形状:", intermediate_products.shape) # (4, 2, 1, 2)
print("中间乘积张量 (部分):\n", intermediate_products[0, 0]) # 示例输出解析 intermediate_products[0, 0] 的含义：
现在，我们得到了一个形状为 (4, 2, 1, 2) 的张量，其中每个元素 intermediate_products[i, j, k, l] 都是 a[i, j, k] * b[j, i, l] 的结果。
为了回到原始的 einsum('ijk,jil->kl', a, b) 行为，我们需要对那些在输出索引字符串 'kl' 中被省略的索引进行求和。在本例中，i 和 j 被省略了。
首先，对索引 j 对应的轴（即 axis=1）进行求和：
# 对 j 轴（axis=1）求和
sum_over_j = intermediate_products.sum(axis=1)
print("对 j 轴求和后的形状:", sum_over_j.shape) # (4, 1, 2)
print("对 j 轴求和后的结果 (部分):\n", sum_over_j[0]) # 示例输出解析 sum_over_j[0] 的含义：
然后，对索引 i 对应的轴（即 axis=0）进行求和：
# 对 i 轴（axis=0）求和
final_result = sum_over_j.sum(axis=0)
print("对 i 轴求和后的形状:", final_result.shape) # (1, 2)
print("最终结果:\n", final_result)解析 final_result 的含义：
这种方法清晰地展示了 einsum 如何先执行元素乘法，然后对未出现在输出索引中的维度进行求和。
另一种深入理解 einsum 的方式是将其转换为等效的显式嵌套循环。这能最直接地展示每个元素是如何被访问、相乘和累加的。
对于 np.einsum('ijk,jil->kl', a, b)，我们可以构建一个等效的 Python 函数：
def sum_array_explicit(A, B):
# 获取张量 A 的形状 (i_len, j_len, k_len)
i_len_A, j_len_A, k_len_A = A.shape
# 获取张量 B 的形状 (j_len, i_len, l_len)
# 注意 B 的索引顺序是 j, i, l，所以这里需要根据 einsum 字符串来理解
# 实际上，einsum 会自动处理索引的匹配和重排
# 我们可以从 einsum 字符串推断出 i, j, k, l 的最大范围
# 假设 A 和 B 的形状是兼容的，我们从 A 和 B 的形状中提取维度长度
# i 对应 A 的第一个维度，j 对应 A 的第二个维度
# k 对应 A 的第三个维度
# j 对应 B 的第一个维度，i 对应 B 的第二个维度
# l 对应 B 的第三个维度
# 确定各个索引的范围
i_max = A.shape[0] # i 的范围由 A 决定
j_max = A.shape[1] # j 的范围由 A 决定
k_max = A.shape[2] # k 的范围由 A 决定
l_max = B.shape[2] # l 的范围由 B 决定 (注意 B 的第三个维度是 l)
# 初始化结果张量，形状为 (k_len, l_len)
ret = np.zeros((k_max, l_max))
# 遍历所有可能的 i, j, k, l 组合
for i in range(i_max):
for j in range(j_max):
for k in range(k_max):
for l in range(l_max):
# 核心操作：A[i, j, k] 乘以 B[j, i, l] 并累加到 ret[k, l]
# 注意 B 的索引顺序是 j, i, l
ret[k, l] += A[i, j, k] * B[j, i, l]
return ret
# 使用示例张量运行显式循环
result_explicit = sum_array_explicit(a, b)
print("显式循环计算结果:\n", result_explicit)
# 与 einsum 的结果进行比较
result_einsum = np.einsum('ijk,jil->kl', a, b)
print("einsum 计算结果:\n", result_einsum)
# 验证结果是否一致
print("两种方法结果是否一致:", np.allclose(result_explicit, result_einsum))关键观察点：
通过这种显式循环的视角，我们可以清晰地看到 einsum 如何在底层执行元素级的乘法和累加操作，以及索引字符串如何精确地指导这些操作。
np.einsum 提供了一种高度灵活和表达性强的张量操作方式。理解其工作原理的关键在于：
使用 einsum 的优势：
注意事项：
通过上述两种方法，我们不仅看到了 np.einsum('ijk,jil->kl', a, b) 的最终结果，更深入地理解了其内部元素如何根据索引字符串进行乘法、重排和求和，从而为更高级的张量操作打下了坚实的基础。
以上就是深入理解 NumPy einsum 的张量求和机制的详细内容，更多请关注php中文网其它相关文章！
每个人都需要一台速度更快、更稳定的 PC。随着时间的推移，垃圾文件、旧注册表数据和不必要的后台进程会占用资源并降低性能。幸运的是，许多工具可以让 Windows 保持平稳运行。
Copyright 2014-2025 https://www.php.cn/ All Rights Reserved | php.cn | 湘ICP备2023035733号
130
收藏
