Repository: dask/dask

Issue:
Better Dask Array support of `frompyfunc`
NumPy's [`frompyfunc`]( https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.frompyfunc.html ) is quite useful for wrapping up a basic Python function as a `ufunc`. It also works pretty much out of the box with Dask Arrays thanks to [`__array_ufunc__`]( https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html#array-ufunc-for-ufuncs ).

However there are two caveats. First one must have NumPy 1.13 at least for `__array_ufunc__` to work (we support 1.10+ currently). Second it appears that `frompyfunc` objects cannot be pickled. ( https://github.com/numpy/numpy/issues/8099 ) The latter hurts any distributed computing use cases where the `frompyfunc` object must be shipped between processes.

One way to fix both of these would be to create some sort of wrapper object derived from [Dask Array's `ufunc`]( https://github.com/dask/dask/blob/9e9fa10ef11bfe86a52214d8a2cda19508a4ee37/dask/array/ufunc.py#L45 ) that delays creating of the `frompyfunc` object until the computation is performed. This should allow the `frompyfunc` object to interact more reliably with Dask Arrays on older NumPy versions. Further it would allow the object to pickled more easily (assuming the Python function used can be pickled).


Discussion:
In general I'm not opposed to continuing to wrap numpy ufuncs.  Presumably
as long as the function itself was serializable this would work fine in a
distributed context.

On Mon, Dec 4, 2017 at 2:54 AM, jakirkham <notifications@github.com> wrote:

> NumPy's frompyfunc
> <https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.frompyfunc.html>
> is quite useful for wrapping up a basic Python function as ufunc. It also
> works pretty much out of the box with Dask Arrays thanks to
> __array_ufunc__
> <https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html#array-ufunc-for-ufuncs>
> .
>
> However there are two caveats. First one must have NumPy 1.13 at least for
> __array_ufunc__ to work (we support 1.10+ currently). Second it appears
> that frompyfunc objects cannot be pickled. ( numpy/numpy#8099
> <https://github.com/numpy/numpy/issues/8099> ) The latter hurts any
> distributed computing use cases where the frompyfunc object must be
> shipped between processes.
>
> One way to fix both of these would be to create some sort of wrapper
> object derived from Dask Array's ufunc
> <https://github.com/dask/dask/blob/9e9fa10ef11bfe86a52214d8a2cda19508a4ee37/dask/array/ufunc.py#L45>
> that delays creating of the frompyfunc object until the computation is
> performed. This should allow the frompyfunc object to interact more
> reliably with Dask Arrays on older NumPy versions. Further it would allow
> the object to pickled more easily (assuming the Python function used can be
> pickled).
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/2956>, or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszDGG_BylJAjH2WS4VeMzB-CS25j6ks5s86U5gaJpZM4Q0P3C>
> .
>

It seems like a lot of user questions end up going along the lines of "I can do this special thing with NumPy how do I take it to larger data with Dask?" Hence there is a lot of value in having operations in Dask that generally wrap up things of this sort (e.g. `map_blocks`, `atop`, `apply_along_axis`, etc.). My thinking with raising `frompyfunc` in particular is it handles another class of operations like this. Namely wrapping up user provided Python functions that are intended to act in an element-wise fashion.

Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -24,7 +24,7 @@
         fmax, fmin, isreal, iscomplex, isfinite, isinf, isnan, signbit,
         copysign, nextafter, spacing, ldexp, fmod, floor, ceil, trunc, degrees,
         radians, rint, fix, angle, real, imag, clip, fabs, sign, absolute,
-        i0, sinc, nan_to_num, frexp, modf, divide)
+        i0, sinc, nan_to_num, frexp, modf, divide, frompyfunc)
 try:
     from .ufunc import float_power
 except ImportError:
diff --git a/dask/array/ufunc.py b/dask/array/ufunc.py
--- a/dask/array/ufunc.py
+++ b/dask/array/ufunc.py
@@ -1,15 +1,15 @@
 from __future__ import absolute_import, division, print_function
 
 from operator import getitem
-from functools import partial
+from functools import partial, wraps
 
 import numpy as np
 from toolz import curry
 
 from .core import Array, elemwise, atop, apply_infer_dtype, asarray
-from ..base import is_dask_collection
+from ..base import is_dask_collection, normalize_function
 from .. import core, sharedict
-from ..utils import skip_doctest
+from ..utils import skip_doctest, funcname
 
 
 def __array_wrap__(numpy_ufunc, x, *args, **kwargs):
@@ -42,14 +42,56 @@ def wrapped(*args, **kwargs):
     return wrapped
 
 
+class da_frompyfunc(object):
+    """A serializable `frompyfunc` object"""
+    def __init__(self, func, nin, nout):
+        self._ufunc = np.frompyfunc(func, nin, nout)
+        self._func = func
+        self.nin = nin
+        self.nout = nout
+        self._name = funcname(func)
+        self.__name__ = 'frompyfunc-%s' % self._name
+
+    def __repr__(self):
+        return 'da.frompyfunc<%s, %d, %d>' % (self._name, self.nin, self.nout)
+
+    def __dask_tokenize__(self):
+        return (normalize_function(self._func), self.nin, self.nout)
+
+    def __reduce__(self):
+        return (da_frompyfunc, (self._func, self.nin, self.nout))
+
+    def __call__(self, *args, **kwargs):
+        return self._ufunc(*args, **kwargs)
+
+    def __getattr__(self, a):
+        if not a.startswith('_'):
+            return getattr(self._ufunc, a)
+        raise AttributeError("%r object has no attribute "
+                             "%r" % (type(self).__name__, a))
+
+    def __dir__(self):
+        o = set(dir(type(self)))
+        o.update(self.__dict__)
+        o.update(dir(self._ufunc))
+        return list(o)
+
+
+@wraps(np.frompyfunc)
+def frompyfunc(func, nin, nout):
+    if nout > 1:
+        raise NotImplementedError("frompyfunc with more than one output")
+    return ufunc(da_frompyfunc(func, nin, nout))
+
+
 class ufunc(object):
     _forward_attrs = {'nin', 'nargs', 'nout', 'ntypes', 'identity',
                       'signature', 'types'}
 
     def __init__(self, ufunc):
-        if not isinstance(ufunc, np.ufunc):
-            raise TypeError("must be an instance of `ufunc`, "
-                            "got `%s" % type(ufunc).__name__)
+        if not isinstance(ufunc, (np.ufunc, da_frompyfunc)):
+            raise TypeError("must be an instance of `ufunc` or "
+                            "`da_frompyfunc`, got `%s" % type(ufunc).__name__)
         self._ufunc = ufunc
         self.__name__ = ufunc.__name__
         copy_docstring(self, ufunc)


Test patch:
diff --git a/dask/array/tests/test_ufunc.py b/dask/array/tests/test_ufunc.py
--- a/dask/array/tests/test_ufunc.py
+++ b/dask/array/tests/test_ufunc.py
@@ -1,12 +1,16 @@
 from __future__ import absolute_import, division, print_function
 
+import pickle
 from functools import partial
+from operator import add
 
 import pytest
 np = pytest.importorskip('numpy')
 
 import dask.array as da
+from dask.array.ufunc import da_frompyfunc
 from dask.array.utils import assert_eq
+from dask.base import tokenize
 
 
 def test_ufunc_meta():
@@ -263,6 +267,56 @@ def test_angle():
     assert_eq(da.angle(comp), np.angle(comp))
 
 
+def test_frompyfunc():
+    myadd = da.frompyfunc(add, 2, 1)
+    np_myadd = np.frompyfunc(add, 2, 1)
+
+    x = np.random.normal(0, 10, size=(10, 10))
+    dx = da.from_array(x, chunks=(3, 4))
+    y = np.random.normal(0, 10, size=10)
+    dy = da.from_array(y, chunks=2)
+
+    assert_eq(myadd(dx, dy), np_myadd(x, y))
+    assert_eq(myadd.outer(dx, dy), np_myadd.outer(x, y))
+
+    with pytest.raises(NotImplementedError):
+        da.frompyfunc(lambda x, y: (x + y, x - y), 2, 2)
+
+
+def test_frompyfunc_wrapper():
+    f = da_frompyfunc(add, 2, 1)
+    np_f = np.frompyfunc(add, 2, 1)
+    x = np.array([1, 2, 3])
+
+    # Callable
+    np.testing.assert_equal(f(x, 1), np_f(x, 1))
+
+    # picklable
+    f2 = pickle.loads(pickle.dumps(f))
+    np.testing.assert_equal(f2(x, 1), np_f(x, 1))
+
+    # Attributes
+    assert f.ntypes == np_f.ntypes
+    with pytest.raises(AttributeError):
+        f.not_an_attribute
+
+    # Tab completion
+    assert 'ntypes' in dir(f)
+
+    # Methods
+    np.testing.assert_equal(f.outer(x, x), np_f.outer(x, x))
+
+    # funcname
+    assert f.__name__ == 'frompyfunc-add'
+
+    # repr
+    assert repr(f) == "da.frompyfunc<add, 2, 1>"
+
+    # tokenize
+    assert (tokenize(da_frompyfunc(add, 2, 1)) ==
+            tokenize(da_frompyfunc(add, 2, 1)))
+
+
 @pytest.mark.skipif(np.__version__ < '1.13.0', reason='array_ufunc not present')
 def test_array_ufunc():
     x = np.arange(24).reshape((4, 6))