我正在研究一个基本上可以归结为求解矩阵方程的项目
A.dot(x) = d
其中A是一个矩阵，其尺寸大约为 10 000 000 到 2000 （我想最终在两个方向上增加它）。
A显然不适合内存，因此必须并行化。我通过解决A.T.dot(A).dot(x) = A.T.dot(d)来做到这一点。A.T到 2000 年的维度将是 2000。它可以通过将A和d分成块来计算，A_i并且d_i沿着行计算A_i.T.dot(A_i)和A_i.T.dot(d_i)，然后将它们相加。非常适合并行化。我已经能够使用多处理模块来实现这一点，但是 1）A由于内存使用，很难进一步扩展（在两个维度上都增加），以及 2）不漂亮（因此不容易维护）。
Dask 似乎是一个非常有前途的库，可以解决这两个问题，我也做了一些尝试。我的A矩阵计算起来很复杂：它基于大约 15 个不同的数组（大小等于 中的行数A），其中一些用于迭代算法来评估相关的勒让德函数。当chunk很小（10000行）时，构建任务图需要很长时间，并且占用大量内存（内存增加与迭代算法的调用不谋而合）。当 chunk 较大（50000 行）时，计算前的内存消耗会小很多，但在计算时会迅速耗尽A.T.dot(A)。我已经尝试过cache.Chest，但它显着减慢了计算速度。
任务图一定很大很复杂——调用A._visualize()崩溃。使用更简单A的矩阵，它可以直接执行此操作（请参阅@MRocklin 的回复）。有没有办法让我简化它？
任何有关如何解决此问题的建议将不胜感激。
作为一个玩具示例，我尝试了
A = da.ones((2e3, 1e7), chunks = (2e3, 1e3))
(A.T.dot(A)).compute()
这也失败了，耗尽了所有内存，只有一个内核处于活动状态。使用chunks = (2e3, 1e5)时，所有内核几乎立即启动，但MemoryError会在 1 秒内出现（我当前的计算机上有 15 GB）。chunks = (2e3, 1e4)更有希望，但它最终也消耗了所有内存。
编辑：我删除了玩具示例测试，因为尺寸错误，并更正了其余的尺寸。正如@MRocklin 所说，它确实适用于正确的尺寸。我添加了一个我现在认为与我的问题更相关的问题。
Edit2：这是我试图做的一个非常简化的例子。我相信，问题在于定义A.
import dask.array as da
N = 1e6
M = 500
x = da.random.random((N, 1), chunks = 5*M)
# my actual
A_dict = {0:x}
for i in range(1, M):
A_dict[i] = 2*A_dict[i-1]
A = da.hstack(tuple(A_dict.values()))
A = A.rechunk((M*5, M))
ATA = A.T.dot(A)
这似乎导致了一个非常复杂的任务图，甚至在计算开始之前就占用了大量内存。
我现在已经解决了这个问题，方法是将递归放在一个带有numpy数组的函数中，或多或少地做A = x.map_blocks(...).
作为第二点，一旦我有了A矩阵任务图，A.T.dot(A)直接计算似乎确实会产生一些内存问题（内存使用情况不是很稳定）。因此，我明确地分块计算它，并对结果求和。即使有这些变通方法，dask 在速度和可读性方面也有很大的不同。