python - 如何(/if)使用 dask 转置分布式 3D numpy 数组？

Question

我的问题是在 3D numpy 数组 A 上执行 3 次矩阵乘法，该数组太大而无法容纳在单个处理器中。在张量形式中，我想要 A_ijk B_km C_jn D_ip (B、C 和 D 都可以放入内存中)。我想知道 dask 是否适合此任务(或者其他工具是否更适合)。

我认为最好的方法是将此操作拆分为每个乘法，并确保它们都是本地的。此链接有一个非常有用的图表，总结了我正在谈论的内容 http://www.2decomp.org/1d_mode.html .

更详细地说:首先，要执行 A_ijk B_km，我应该将 A 分布在前两个轴上，并在本地对每支铅笔执行矩阵乘法(图中的第一步)。

然后，我需要转置数组，使 j 轴位于每个处理器的本地(并在 k(现在是 m)轴上分割)，然后执行下一个乘法。 (因此从图中的第一步到第二步)。这就是我想知道 dask 是否可以提供帮助的地方。

我知道原则上可以使用 mpi4py 来完成此操作，但步骤非常重要，而 dask 数组具有有用的重新分块和转置方法，感觉与此应用程序相关。

这看起来是不是很适合 dask？

如果没有，是否有人知道任何可以执行这些步骤的 python 库？我知道fftw有执行此操作的例程，但我不知道如何编写必要的 C 代码，或者如何让它与 python 和 numpy 交互。

感谢您的帮助。

Answer 1

对于 future 的其他人来说，mpi4py 确实有一个转置方法。但它被称为 Alltoall/Alltoallv。 mpi4py 的文档或教程中没有对此进行解释。我在另一个教程中发现了它:https://info.gwdg.de/wiki/doku.php?id=wiki:hpc:mpi4py .