`np.einsum` 提供了一种强大且简洁的方式来执行张量乘法和求和。本文旨在深入剖析 `einsum` 在处理两个张量（例如 `’ijk,jil->kl’`）时，其内部元素是如何进行乘法和求和的。我们将通过分解中间步骤和构建等效的显式循环两种方法，详细揭示 `einsum` 基于索引符号实现复杂张量运算的内在机制，帮助读者全面掌握其工作原理。
NumPy 中的 einsum 函数以其简洁而强大的张量操作能力而闻名，它允许用户通过爱因斯坦求和约定来表达复杂的张量运算，包括点积、外积、转置、张量收缩等。然而，对于初学者而言，理解 einsum 在执行如 np.einsum(‘ijk,jil->kl’, a, b) 这类操作时，其内部元素是如何进行组合和求和的，可能是一个挑战。本文将通过具体的示例和两种不同的解析方法，详细阐述这一过程。
假设我们有两个 NumPy 张量 a 和 b：
import numpy as npa = 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’ 指示了以下操作：
a 张量的维度由 i, j, k 索引。b 张量的维度由 j, i, l 索引。结果张量的维度由 k, l 索引。
这意味着：
i 和 j 是在两个输入张量中都出现的索引。k 和 l 是在结果张量中出现的索引。i 和 j 在输出中被省略，因此它们将是求和的维度。
方法一：通过中间输出分解求和过程
为了理解元素是如何组合和求和的，我们可以首先修改 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] 的含义：
i=0, j=0 时的所有 k, l 组合的乘积。a[0, 0, k] 乘以 b[0, 0, l]。a[0, 0, 0] 是 0.0。b[0, 0, 0] 是 0.0，b[0, 0, 1] 是 1.0。所以 intermediate_products[0, 0, 0, 0] = a[0, 0, 0] * b[0, 0, 0] = 0.0 * 0.0 = 0.0。intermediate_products[0, 0, 0, 1] = a[0, 0, 0] * b[0, 0, 1] = 0.0 * 1.0 = 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=0 时，对所有 j 的 k, l 组合进行求和。sum_over_j[0, k, l] = intermediate_products[0, 0, k, l] + intermediate_products[0, 1, k, l]。例如，sum_over_j[0, 0, 0] = intermediate_products[0, 0, 0, 0] + intermediate_products[0, 1, 0, 0]= (a[0, 0, 0] * b[0, 0, 0]) + (a[0, 1, 0] * b[1, 0, 0])= (0.0 * 0.0) + (1.0 * 8.0) = 8.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 的含义：
final_result[k, l] 是对所有 i 和 j 的 a[i, j, k] * b[j, i, l] 求和的结果。例如，final_result[0, 0] = sum_over_j[0, 0, 0] + sum_over_j[1, 0, 0] + sum_over_j[2, 0, 0] + sum_over_j[3, 0, 0]。它等价于 np.einsum(‘ijk,jil->kl’, a, b) 的结果。
这种方法清晰地展示了 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))
关键观察点：
循环中的核心行 ret[k, l] += A[i, j, k] * B[j, i, l] 直接反映了 einsum 的索引字符串 ‘ijk,jil->kl’。A[i, j, k] 对应 ijk。B[j, i, l] 对应 jil。ret[k, l] 对应 kl。i 和 j 是在两个输入张量中都出现，但在输出张量中被省略的索引。在显式循环中，它们作为外层循环变量，它们的每次迭代都产生一个乘积，这些乘积最终累加到 ret[k, l] 中，从而实现了对 i 和 j 的求和。k 和 l 是输出张量的索引，它们决定了 ret 的形状和每个元素的最终值。
通过这种显式循环的视角，我们可以清晰地看到 einsum 如何在底层执行元素级的乘法和累加操作，以及索引字符串如何精确地指导这些操作。
总结与注意事项
np.einsum 提供了一种高度灵活和表达性强的张量操作方式。理解其工作原理的关键在于：
索引匹配与乘法： einsum 会根据输入张量的索引字符串，将相应维度的元素进行匹配和逐元素乘法。例如，’ijk,jil’ 表示 a 的 i,j,k 维度与 b 的 j,i,l 维度进行匹配。求和规则： 任何在输入索引字符串中出现，但在输出索引字符串中被省略的索引，都将被视为求和维度。einsum 会沿着这些维度对乘积结果进行累加。输出维度： 输出索引字符串定义了结果张量的维度顺序和形状。
使用 einsum 的优势：
简洁性： 用一行代码表达复杂的张量运算。灵活性： 几乎可以实现任何线性代数运算。性能： 通常比手动编写的循环更快，因为它在底层利用了优化的 C 或 Fortran 代码。
注意事项：
索引命名： 索引名称本身不重要，但其位置和在输入/输出字符串中的出现情况至关重要。维度兼容性： einsum 会自动检查维度兼容性，例如，如果 a 的 j 维度长度为 2，b 的 j 维度长度为 3，则会报错。理解爱因斯坦求和约定： 掌握这一数学约定是高效使用 einsum 的基础。
通过上述两种方法，我们不仅看到了 np.einsum(‘ijk,jil->kl’, a, b) 的最终结果，更深入地理解了其内部元素如何根据索引字符串进行乘法、重排和求和，从而为更高级的张量操作打下了坚实的基础。
以上就是深入理解 NumPy einsum 的张量求和机制的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1378227.html
微信扫一扫 
支付宝扫一扫 
