python - 如何（/如果）使用 dask 转置分布式 3D numpy 数组？

我的问题是在 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 交互。

谢谢你的帮助。

标签： pythonnumpympidask