如何将Numpy einsum操作转换为Eigen Tensor运算（适配Jet类型）

理解原Einsum操作 #

原Numpy代码中的np.einsum('ijl,jkl->ikl', U, L)等价于：对每个索引l，执行矩阵乘法U[:, :, l] @ L[:, :, l]，最终将所有结果沿l维度堆叠，输出形状为(2,2,136)的张量。

Eigen Tensor实现方案 #

利用Eigen的表达式模板（避免显式循环），通过shuffle（维度重排）和contract（张量收缩）实现，确保Jet类型能被自动向量化优化。

步骤1：定义类型与张量 #

#include <Eigen/Dense>
#include <Eigen/Tensor>

// 假设Jet类型为自动微分的对偶数类型
using Jet = Eigen::AutoDiffScalar<Eigen::VectorXd>;

// 对应原Numpy的维度：I=2, J=2, L_dim=136
constexpr int I = 2, J = 2, L_dim = 136;

// 初始化张量（与Numpy行优先存储一致，指定RowMajor）
Eigen::Tensor<Jet, 3, Eigen::RowMajor> U(I, J, L_dim);
Eigen::Tensor<Jet, 3, Eigen::RowMajor> L(J, I, L_dim); // K=2，对应原L的J,K,L维度

// 省略U和L的初始化代码（对应np.random.rand）

步骤2：通过批量矩阵乘法实现 #

将l维度转为批量维度，执行批量矩阵乘法后再恢复原维度顺序：

// 1. 将U的维度从(I,J,L)重排为(L,I,J)（批量维度前置）
Eigen::array<int, 3> u_perm = {2, 0, 1};
auto U_batch = U.shuffle(u_perm);

// 2. 将L的维度从(J,K,L)重排为(L,J,K)
Eigen::array<int, 3> l_perm = {2, 0, 1};
auto L_batch = L.shuffle(l_perm);

// 3. 执行批量矩阵乘法：每个批量（l）对应矩阵I×J 乘 J×K，得到(L,I,K)
Eigen::array<Eigen::IndexPair<int>, 1> matmul_dims = {Eigen::IndexPair<int>(2, 1)};
auto result_batch = U_batch.contract(L_batch, matmul_dims);

// 4. 将结果维度从(L,I,K)重排为(I,K,L)，匹配原Einsum输出
Eigen::array<int, 3> res_perm = {1, 2, 0};
Eigen::Tensor<Jet, 3, Eigen::RowMajor> result = result_batch.shuffle(res_perm);

步骤3：直接张量收缩实现 #

也可以直接通过张量收缩指定维度映射，一步得到结果：

// 指定收缩维度：U的第1维（J）与L的第0维（J）
Eigen::array<Eigen::IndexPair<int>, 1> contract_dims = {Eigen::IndexPair<int>(1, 0)};

// 直接收缩并指定结果维度顺序：I, K, L
Eigen::Tensor<Jet, 3, Eigen::RowMajor> result = U.contract(L, contract_dims).reshape(Eigen::array<int, 3>{I, I, L_dim});

关键优势 #

• 所有操作（shuffle/contract/reshape）均为Eigen表达式模板，不会立即执行，编译器可将整个操作融合为向量化内核。
• 完全兼容Jet类型，Eigen会自动重载操作符实现对数值和梯度数组的向量化计算，无需手动循环。

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