Моя задача состоит в том, чтобы выполнить 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)), чтобы затем выполнить следующее умножение. (Итак, переходя от первого ко второму шагу на диаграмме). Вот где мне интересно, может ли Даск помочь.
Я знаю, что в принципе это можно сделать с помощью mpi4py, но шаги довольно нетривиальны, в то время как массивы dask имеют полезные методы rechunk и транспонирования, которые кажутся актуальными для этого приложения.
Кажется ли это чем-то хорошо подходящим для темных?
Если нет, кто-нибудь знает какие-либо библиотеки Python, которые могут выполнять эти шаги? Я знаю, что fftw имеет подпрограммы для выполнения именно этого, но я не знаю, как написать необходимый C-код или как заставить его взаимодействовать с python и numpy.
Спасибо за любую помощь.