Numpy 3D Array/Eigen 张量 C++，将每一行乘以其转置

Question

假设我有一个形状为 (55, 50, 2) 的 numpy 数组 A，我想执行以下操作

B = np.dot(A[0, :, 0][:, None], A[0, :, 0][None, :])

即通过转置计算每行的点积（在 i 和 k 上）

如果不使用 np.einsum 并且显然没有任何 for 循环，如何通过纯广播和重塑（如果需要）来完成此操作？

注：

我在这里标记 eigen3 因为基本上我想用 Eigen::Tensor 根据 .broadcast() 和 .reshape() 重写这个操作（我写起来更容易在 numpy 中先出来以获取一般图片）。因此，直接 Eigen 解决方案将不胜感激，但由于 python，因此 numpy 更受欢迎，我也会接受 numpy 解决方案。

Answer 1

所以你需要的是每一行的外积。

如果使用 C++ Eigen，最好编写循环并进行矩阵乘法。您可以派遣到 BLAS 进行性能测试。

如果你想要一种使用 Numpy 的方法，一个混乱的解决方案是

B = A[...,newaxis].transpose(0,2,1,3)
C = B@B.transpose(0,1,3,2) #Matrix multiplication
C = C.transpose(0,3,2,1)
np.array_equal(np.einsum('ijk,ilk->ijlk',A,A),C) ## check if both are identical

Answer 2

如果我没理解错的话，那么这个操作可以使用Eigen::Tensor模块来完成，如下

#include <unsupported/Eigen/CXX11/Tensor>
int main(){
    // Define a tensor
    Eigen::Tensor<double,3> A(55,50,2);
    A.setRandom();

    // Define the pairs of indices to multiply (i.e. index 1 on A twice)
    std::array<Eigen::IndexPair<long>, 1> idx = { Eigen::IndexPair<long> {1, 1} };
    
    // Contract. B will have dimensions [55,2,55,2]
    Eigen::Tensor<double,4> B = A.contract(A, idx);
}

See on compiler-explorer.

并行执行可以通过一个小的修改来完成：

#include <unsupported/Eigen/CXX11/Tensor>
int main(){
    // Define a tensor
    Eigen::Tensor<double,3> A(55,50,2);
    A.setRandom();

    // Define the pairs of indices to multiply (i.e. index 1 on A twice)
    std::array<Eigen::IndexPair<long>, 1> idx = { Eigen::IndexPair<long> {1, 1} };
    
    // Define a parallel device 
    int num_threads = 4;
    Eigen::ThreadPool       tpl (num_threads);
    Eigen::ThreadPoolDevice dev (&tpl, num_threads);

    // Initialize result buffer B. B will have dimensions [55,2,55,2]
    Eigen::Tensor<double,4> B(55,2,55,2);

    // Contract using the parallel device
    B.device(dev) = A.contract(A, idx);
}

See on compiler-explorer。注意编译器标志，特别是 -DEIGEN_USE_THREADS。另外，我必须设置 -O0 因为 -O3 超时。

关于评论中讨论的性能主题：根据我过去 ~4 年每天使用它的经验，使用 Eigen::Tensor 模块的收缩效果非常好，我无法 find/produce 任何更快的东西（即共享内存上密集的、动态大小的张量）。根据this answer (by the original author of Eigen), the contractions fall back to the matrix product kernels when possible, which should give some confidence. In this case you can use any BLAS library as backend.

Numpy 3D Array/Eigen 张量 C++，将每一行乘以其转置

Numpy 3D Array/Eigen Tensor C++, multiply each row by its transpose

numpy

eigen