dask.array.apply_gufunc с несколькими выходами разных форм

#dask #python-xarray

Вопрос:

Я пытаюсь применить ufunc к фрагментированным широковещательным массивам dask, которые выдают несколько выходов разных форм:

 import dask.array as da # dask.__version__ is 1.2.0
import numpy as np

def func(A3, A2):
    return A3 A2, A2**2

A3 = da.from_array(np.random.randn(3,5,5), chunks=(3,2,2))
A2 = da.from_array(np.random.randn(  5,5), chunks=(  2,2))
ret = da.apply_gufunc(func, '(),()->(),()', A3, A2, output_dtypes=[float,float])

for r in ret:
    print(r)
    r.compute()

Проблема в том, что предполагается, что оба вывода in ret имеют форму (3,5,5), которая затем .compute() завершается ValueError: axes don't match array ошибкой для второго вывода, который должен быть 2-мерным.

Как я могу использовать apply_gufunc в этом случае?

Примечание: я бы, вероятно, предпочел использовать xarray.apply_ufunc в этом случае, но, к сожалению, он пока недоступен с несколькими выходами (см. Здесь ).

1. Да, получил значок «перекати-поле» для этого вопроса o/

2. К сожалению, не так много сопровождающих, которые знакомы с семантикой gufunc @shoyer — ваш лучший выбор, но я думаю, что он довольно занят в эти дни.

Ответ №1:

Вот обходной путь, который помогает:

 def func(A3, A2):
    return A3 A2, (A2**2)[np.newaxis,:]

1. Спасибо за обходной путь! Проблема в том, что все выходные данные приводятся к одной и той же форме, таким образом ret[1].shape == (3, 5, 5) . Хотелось бы, чтобы он оставался в форме (5, 5) или, по крайней мере (1, 5, 5) ,.