深入理解 NumPy einsum 操作的细节与机制

import numpy as np

a = np.arange(8.).reshape(4, 2, 1) # 形状 (4, 2, 1)
b = np.arange(16.).reshape(2, 4, 2) # 形状 (2, 4, 2)

print("Tensor a:\n", a)
print("Tensor b:\n", b)

# 步骤1: 获取所有未经求和的乘积
products = np.einsum('ijk,jil->ijkl', a, b)
print("所有未经求和的乘积 (shape:", products.shape, "):\n", products)

# 步骤2: 对 j 轴（products 的第1轴）进行求和
sum_over_j = products.sum(axis=1)
print("\n对 j 轴求和后 (shape:", sum_over_j.shape, "):\n", sum_over_j)

# 步骤3: 对 i 轴（sum_over_j 的第0轴）进行求和
final_result = sum_over_j.sum(axis=0)
print("\n对 i 轴求和后 (shape:", final_result.shape, "):\n", final_result)

# 验证与原始 einsum 结果一致
original_einsum_result = np.einsum('ijk,jil->kl', a, b)
print("\n原始 einsum 结果 (shape:", original_einsum_result.shape, "):\n", original_einsum_result)

assert np.allclose(final_result, original_einsum_result)
print("\n分解求和结果与原始 einsum 结果一致。")

def sum_array_explicit_loop(A, B):
    # 获取张量 A 的维度长度
    i_len, j_len, k_len = A.shape
    # 获取张量 B 的维度长度 (注意 B 的形状是 (j_len, i_len, l_len) 
    # 如果按照 einsum 的 jil 索引来理解，但其原始形状是 (2, 4, 2)，
    # 这里的 _ 和 l_len 对应 B 的第0维和第2维)
    # 实际上，B 的原始形状是 (B_dim0, B_dim1, B_dim2)
    # 在 'jil' 中，j 对应 B_dim0, i 对应 B_dim1, l 对应 B_dim2
    # 所以，B.shape[0] 是 j 的最大值，B.shape[1] 是 i 的最大值，B.shape[2] 是 l 的最大值
    # 但是，i_len 和 j_len 已经由 A 决定，所以我们只需要 l_len
    # 确保维度兼容性：A.shape[1] (j_len_A) 必须等于 B.shape[0] (j_len_B)
    # A.shape[0] (i_len_A) 必须等于 B.shape[1] (i_len_B)
    # 这里我们直接从 A 和 B 的实际形状推导循环范围

    # 重新确认循环范围的正确性：
    # i 循环范围由 A.shape[0] 决定
    # j 循环范围由 A.shape[1] 决定
    # k 循环范围由 A.shape[2] 决定
    # l 循环范围由 B.shape[2] 决定 (因为 B 的第三个索引是 l)

    # 对于 'ijk,jil->kl'
    # i 的范围是 A.shape[0]
    # j 的范围是 A.shape[1] (同时也是 B.shape[0])
    # k 的范围是 A.shape[2]
    # l 的范围是 B.shape[2]

    i_max = A.shape[0]
    j_max = A.shape[1]
    k_max = A.shape[2]
    l_max = B.shape[2] # l 是 B 的最后一个维度

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

    # 四重嵌套循环模拟 einsum 运算
    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，这意味着 B 的原始第0维对应 j，第1维对应 i，第2维对应 l
                    ret[k, l] += A[i, j, k] * B[j, i, l]
    return ret

# 使用显式循环计算结果
explicit_loop_result = sum_array_explicit_loop(a, b)
print("\n显式循环计算结果 (shape:", explicit_loop_result.shape, "):\n", explicit_loop_result)

assert np.allclose(explicit_loop_result, original_einsum_result)
print("\n显式循环结果与原始 einsum 结果一致。")