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

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

3 голосов
/ 17 апреля 2019

Я пытаюсь применить 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()

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

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

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

1 Ответ

1 голос
/ 07 июня 2019

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

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