寻求与einsum全约简等价的N维tensordot通用实现方案

嘿，这个需求我太懂了！每次针对不同维度写einsum确实麻烦，用tensordot完全可以实现通用的任意N阶张量全约简，不用逐个维度编写代码。

核心思路 #

你用einsum('ijk,ijk', A, B)做的事情，本质是两个同形状张量对应位置元素相乘后，对所有元素求和。要让tensordot实现这个效果，只需要告诉它：把两个张量的所有对应轴都进行收缩（也就是一一配对求和）。

通用实现方法 #

对于任意同形状的N阶张量A和B，只需要获取它们的维度数ndim，然后给tensordot的axes参数传入两个包含所有轴序号的列表，或者直接传入维度数ndim（更简洁）：

代码示例

import numpy as np

# 通用函数封装
def full_tensor_reduce(A, B):
    if A.shape != B.shape:
        raise ValueError("两个张量必须具有相同的形状！")
    ndim = A.ndim
    # 两种写法等价，选你喜欢的就行
    # 写法1：显式指定所有轴
    # return np.tensordot(A, B, axes=(list(range(ndim)), list(range(ndim))))
    # 写法2：直接传入维度数，更简洁
    return np.tensordot(A, B, axes=ndim)

验证等价性 #

我们可以用你给出的1D到4D例子来验证，结果和einsum完全一致：

# 1D测试
A_1D = np.random.rand(5)
B_1D = np.random.rand(5)
print(np.allclose(np.einsum('i,i', A_1D, B_1D), full_tensor_reduce(A_1D, B_1D)))  # True

# 2D测试
A_2D = np.random.rand(3,4)
B_2D = np.random.rand(3,4)
print(np.allclose(np.einsum('ij,ij', A_2D, B_2D), full_tensor_reduce(A_2D, B_2D)))  # True

# 3D测试
A_3D = np.random.rand(2,3,4)
B_3D = np.random.rand(2,3,4)
print(np.allclose(np.einsum('ijk,ijk', A_3D, B_3D), full_tensor_reduce(A_3D, B_3D)))  # True

# 4D测试
A_4D = np.random.rand(2,2,3,4)
B_4D = np.random.rand(2,2,3,4)
print(np.allclose(np.einsum('ijkl,ijkl', A_4D, B_4D), full_tensor_reduce(A_4D, B_4D)))  # True

为什么这样可行？ #

tensordot的axes参数有几种用法：

• 当传入整数n时，表示将两个张量的前n个轴一一配对收缩；
• 当传入两个列表(axes_a, axes_b)时，将axes_a中的轴和axes_b中的轴一一配对收缩。

对于N阶张量，不管用哪种写法，都是把两个张量的所有轴一一对应收缩，最终效果就是所有对应位置元素相乘后求和，和你用einsum写的全约简完全等价，而且不管N多大都能通用。

内容的提问来源于stack exchange，提问作者Saxodrum