深入理解 NumPy einsum 的张量运算细节

import numpy as np

# 张量 a 的形状为 (4, 2, 1)
a = np.arange(8.).reshape(4, 2, 1)
print("张量 a:\n", a)
# 张量 b 的形状为 (2, 4, 2)
b = np.arange(16.).reshape(2, 4, 2)
print("\n张量 b:\n", b)

# 步骤 1: 执行元素乘法，不进行任何求和
# 通过在输出下标中包含所有输入下标 (i, j, k, l)，可以查看每个元素的乘积。
# 此时输出张量的形状为 (i_len, j_len, k_len, l_len) = (4, 2, 1, 2)
intermediate_product = np.einsum('ijk,jil->ijkl', a, b)
print("中间乘积 (ijkl):\n", intermediate_product)

# 步骤 2: 对 'j' 维度进行求和
# 原始表达式 'ijk,jil->kl' 中，j 是一个被求和的维度。
# 对应到 intermediate_product，j 是其第二个轴 (axis=1)。
sum_over_j = intermediate_product.sum(axis=1)
print("\n对 'j' 维度求和后的结果:\n", sum_over_j)
# 此时张量形状变为 (i_len, k_len, l_len) = (4, 1, 2)

# 步骤 3: 对 'i' 维度进行求和
# 原始表达式 'ijk,jil->kl' 中，i 也是一个被求和的维度。
# 对应到 sum_over_j，i 是其第一个轴 (axis=0)。
final_result_method1 = sum_over_j.sum(axis=0)
print("\n对 'i' 维度求和后的最终结果:\n", final_result_method1)
# 最终张量形状为 (k_len, l_len) = (1, 2)，与 'kl' 匹配。

def sum_array_explicit_loop(A, B):
    # 获取张量 A 的维度长度
    i_len, j_len, k_len = A.shape
    # 获取张量 B 的最后一个维度长度（对应 l）
    # 注意：B 的实际形状是 (j_max, i_max, l_max)，因此 l_len 对应 B.shape[2]
    l_len = B.shape[2]

    # 初始化结果张量，其形状应为 (k_len, l_len)
    ret = np.zeros((k_len, l_len))

    # 遍历所有可能的 i, j, k, l 组合
    # i 遍历 A 的第一个维度 (0-3)
    # j 遍历 A 的第二个维度 (0-1)
    # k 遍历 A 的第三个维度 (0-0)
    # l 遍历 B 的第三个维度 (0-1)
    for i in range(i_len):
        for j in range(j_len):
            for k in range(k_len):
                for l in range(l_len):
                    # 根据 einsum 表达式 'ijk,jil->kl'，执行元素乘法
                    # A 的索引为 (i, j, k)
                    # B 的索引为 (j, i, l)
                    # 将乘积累加到结果张量 ret[k, l] 中
                    ret[k, l] += A[i, j, k] * B[j, i, l]
    return ret

final_result_method2 = sum_array_explicit_loop(a, b)
print("\n通过显式循环计算的最终结果:\n", final_result_method2)

# 验证与直接使用 einsum 的结果是否一致
einsum_direct_result = np.einsum('ijk,jil->kl', a, b)
print("\n直接使用 einsum 的结果:\n", einsum_direct_result)

# 比较两种方法的结果
print("\n两种方法结果是否一致 (方法1 vs 直接einsum):", np.allclose(final_result_method1, einsum_direct_result))
print("两种方法结果是否一致 (方法2 vs 直接einsum):", np.allclose(final_result_method2, einsum_direct_result))