如何（/如果）使用 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 数组具有有用的 rechunk 和 transpose 方法，感觉与此应用程序相关。

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

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

感谢您的帮助。

Answer 1

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

Answer 2

Dask 实现了einsum，这可能是你想要的，当然还有matmul，如果你想写出操作。只要您的大矩阵 A 是一个 Dask 数组，并且具有合理的块大小，Dask 就会在不耗尽内存的情况下打包您的工作。