4

I am trying to apply a ufunc to chunked broadcastable dask arrays which produce several outputs of different shapes:

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()

The problem is that both outputs in ret are assumed to be of shape (3,5,5), which then fails on .compute() with ValueError: axes don't match array for the second output, which should be 2-dimensional.

How can I use apply_gufunc in this case?

Note: I would probably rather use xarray.apply_ufunc in this case but unfortunately it is not available yet with multiple outputs (see here).

François
  • 7,988
  • 2
  • 21
  • 17
  • Yeah, got the "tumbleweed" badge for that question \o/ – François Apr 24 '19 at 20:12
  • Unfortunately there are not very many maintainers who are familiar with gufunc semantics @shoyer is your best bet, but I think he's quite busy these days. – MRocklin Apr 27 '19 at 15:13

1 Answers1

1

Here is a workaround that helps:

def func(A3, A2):
    return A3+A2, (A2**2)[np.newaxis,:]
Ales
  • 495
  • 3
  • 11
  • Thanks for the workaround! The problem is that all outputs get coerced to the same shape, thus `ret[1].shape == (3, 5, 5)`. It would like it to remain of shape `(5, 5)`, or at least `(1, 5, 5)`. – François Jun 08 '19 at 14:27