[Numpy-discussion] Revised NEP-18, __array_function__ protocol

Marten van Kerkwijk m.h.vankerkwijk at gmail.com
Wed Jun 27 11:41:39 EDT 2018

Hi Hameer,

I'm confused: Isn't your reference array just `self`?
All the best,


On Wed, Jun 27, 2018 at 2:27 AM, Hameer Abbasi <einstein.edison at gmail.com>

> On 27. Jun 2018 at 07:48, Stephan Hoyer <shoyer at gmail.com> wrote:
> After much discussion (and the addition of three new co-authors!), I’m
> pleased to present a significantly revision of NumPy Enhancement Proposal
> 18: A dispatch mechanism for NumPy's high level array functions:
> http://www.numpy.org/neps/nep-0018-array-function-protocol.html
> The full text is also included below.
> Best,
> Stephan
> ===========================================================
> A dispatch mechanism for NumPy's high level array functions
> ===========================================================
> :Author: Stephan Hoyer <shoyer at google.com>
> :Author: Matthew Rocklin <mrocklin at gmail.com>
> :Author: Marten van Kerkwijk <mhvk at astro.utoronto.ca>
> :Author: Hameer Abbasi <hameerabbasi at yahoo.com>
> :Author: Eric Wieser <wieser.eric at gmail.com>
> :Status: Draft
> :Type: Standards Track
> :Created: 2018-05-29
> Abstact
> -------
> We propose the ``__array_function__`` protocol, to allow arguments of NumPy
> functions to define how that function operates on them. This will allow
> using NumPy as a high level API for efficient multi-dimensional array
> operations, even with array implementations that differ greatly from
> ``numpy.ndarray``.
> Detailed description
> --------------------
> NumPy's high level ndarray API has been implemented several times
> outside of NumPy itself for different architectures, such as for GPU
> arrays (CuPy), Sparse arrays (scipy.sparse, pydata/sparse) and parallel
> arrays (Dask array) as well as various NumPy-like implementations in the
> deep learning frameworks, like TensorFlow and PyTorch.
> Similarly there are many projects that build on top of the NumPy API
> for labeled and indexed arrays (XArray), automatic differentiation
> (Autograd, Tangent), masked arrays (numpy.ma), physical units
> (astropy.units,
> pint, unyt), etc. that add additional functionality on top of the NumPy
> API.
> Most of these project also implement a close variation of NumPy's level
> high
> API.
> We would like to be able to use these libraries together, for example we
> would like to be able to place a CuPy array within XArray, or perform
> automatic differentiation on Dask array code. This would be easier to
> accomplish if code written for NumPy ndarrays could also be used by
> other NumPy-like projects.
> For example, we would like for the following code example to work
> equally well with any NumPy-like array object:
> .. code:: python
>     def f(x):
>         y = np.tensordot(x, x.T)
>         return np.mean(np.exp(y))
> Some of this is possible today with various protocol mechanisms within
> NumPy.
> -  The ``np.exp`` function checks the ``__array_ufunc__`` protocol
> -  The ``.T`` method works using Python's method dispatch
> -  The ``np.mean`` function explicitly checks for a ``.mean`` method on
>    the argument
> However other functions, like ``np.tensordot`` do not dispatch, and
> instead are likely to coerce to a NumPy array (using the ``__array__``)
> protocol, or err outright. To achieve enough coverage of the NumPy API
> to support downstream projects like XArray and autograd we want to
> support *almost all* functions within NumPy, which calls for a more
> reaching protocol than just ``__array_ufunc__``. We would like a
> protocol that allows arguments of a NumPy function to take control and
> divert execution to another function (for example a GPU or parallel
> implementation) in a way that is safe and consistent across projects.
> Implementation
> --------------
> We propose adding support for a new protocol in NumPy,
> ``__array_function__``.
> This protocol is intended to be a catch-all for NumPy functionality that
> is not covered by the ``__array_ufunc__`` protocol for universal functions
> (like ``np.exp``). The semantics are very similar to ``__array_ufunc__``,
> except
> the operation is specified by an arbitrary callable object rather than a
> ufunc
> instance and method.
> A prototype implementation can be found in
> `this notebook <https://nbviewer.jupyter.org/gist/shoyer/
> 1f0a308a06cd96df20879a1ddb8f0006>`_.
> The interface
> ~~~~~~~~~~~~~
> We propose the following signature for implementations of
> ``__array_function__``:
> .. code-block:: python
>     def __array_function__(self, func, types, args, kwargs)
> -  ``func`` is an arbitrary callable exposed by NumPy's public API,
>    which was called in the form ``func(*args, **kwargs)``.
> -  ``types`` is a ``frozenset`` of unique argument types from the original
> NumPy
>    function call that implement ``__array_function__``.
> -  The tuple ``args`` and dict ``kwargs`` are directly passed on from the
>    original call.
> Unlike ``__array_ufunc__``, there are no high-level guarantees about the
> type of ``func``, or about which of ``args`` and ``kwargs`` may contain
> objects
> implementing the array API.
> As a convenience for ``__array_function__`` implementors, ``types``
> provides all
> argument types with an ``'__array_function__'`` attribute. This
> allows downstream implementations to quickly determine if they are likely
> able
> to support the operation. A ``frozenset`` is used to ensure that
> ``__array_function__`` implementations cannot rely on the iteration order
> of
> ``types``, which would facilitate violating the well-defined "Type casting
> hierarchy" described in
> `NEP-13 <https://www.numpy.org/neps/nep-0013-ufunc-overrides.html>`_.
> Example for a project implementing the NumPy API
> ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
> Most implementations of ``__array_function__`` will start with two
> checks:
> 1.  Is the given function something that we know how to overload?
> 2.  Are all arguments of a type that we know how to handle?
> If these conditions hold, ``__array_function__`` should return
> the result from calling its implementation for ``func(*args, **kwargs)``.
> Otherwise, it should return the sentinel value ``NotImplemented``,
> indicating
> that the function is not implemented by these types. This is preferable to
> raising ``TypeError`` directly, because it gives *other* arguments the
> opportunity to define the operations.
> There are no general requirements on the return value from
> ``__array_function__``, although most sensible implementations should
> probably
> return array(s) with the same type as one of the function's arguments.
> If/when Python gains
> `typing support for protocols <https://www.python.org/dev/peps/pep-0544/
> >`_
> and NumPy adds static type annotations, the ``@overload`` implementation
> for ``SupportsArrayFunction`` will indicate a return type of ``Any``.
> It may also be convenient to define a custom decorators (``implements``
> below)
> for registering ``__array_function__`` implementations.
> .. code:: python
>     class MyArray:
>         def __array_function__(self, func, types, args, kwargs):
>             if func not in HANDLED_FUNCTIONS:
>                 return NotImplemented
>             # Note: this allows subclasses that don't override
>             # __array_function__ to handle MyArray objects
>             if not all(issubclass(t, MyArray) for t in types):
>                 return NotImplemented
>             return HANDLED_FUNCTIONS[func](*args, **kwargs)
>     def implements(numpy_function):
>         """Register an __array_function__ implementation for MyArray
> objects."""
>         def decorator(func):
>             HANDLED_FUNCTIONS[numpy_function] = func
>             return func
>         return decorator
>     @implements(np.concatenate)
>     def concatenate(arrays, axis=0, out=None):
>         ...  # implementation of concatenate for MyArray objects
>     @implements(np.broadcast_to)
>     def broadcast_to(array, shape):
>         ...  # implementation of broadcast_to for MyArray objects
> Note that it is not required for ``__array_function__`` implementations to
> include *all* of the corresponding NumPy function's optional arguments
> (e.g., ``broadcast_to`` above omits the irrelevant ``subok`` argument).
> Optional arguments are only passed in to ``__array_function__`` if they
> were explicitly used in the NumPy function call.
> Necessary changes within the NumPy codebase itself
> ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
> This will require two changes within the NumPy codebase:
> 1. A function to inspect available inputs, look for the
>    ``__array_function__`` attribute on those inputs, and call those
>    methods appropriately until one succeeds.  This needs to be fast in the
>    common all-NumPy case, and have acceptable performance (no worse than
>    linear time) even if the number of overloaded inputs is large (e.g.,
>    as might be the case for `np.concatenate`).
>    This is one additional function of moderate complexity.
> 2. Calling this function within all relevant NumPy functions.
>    This affects many parts of the NumPy codebase, although with very low
>    complexity.
> Finding and calling the right ``__array_function__``
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> Given a NumPy function, ``*args`` and ``**kwargs`` inputs, we need to
> search through ``*args`` and ``**kwargs`` for all appropriate inputs
> that might have the ``__array_function__`` attribute. Then we need to
> select among those possible methods and execute the right one.
> Negotiating between several possible implementations can be complex.
> Finding arguments
> '''''''''''''''''
> Valid arguments may be directly in the ``*args`` and ``**kwargs``, such
> as in the case for ``np.tensordot(left, right, out=out)``, or they may
> be nested within lists or dictionaries, such as in the case of
> ``np.concatenate([x, y, z])``. This can be problematic for two reasons:
> 1. Some functions are given long lists of values, and traversing them
>    might be prohibitively expensive.
> 2. Some functions may have arguments that we don't want to inspect, even
>    if they have the ``__array_function__`` method.
> To resolve these issues, NumPy functions should explicitly indicate which
> of their arguments may be overloaded, and how these arguments should be
> checked. As a rule, this should include all arguments documented as either
> ``array_like`` or ``ndarray``.
> We propose to do so by writing "dispatcher" functions for each overloaded
> NumPy function:
> - These functions will be called with the exact same arguments that were
> passed
>   into the NumPy function (i.e., ``dispatcher(*args, **kwargs)``), and
> should
>   return an iterable of arguments to check for overrides.
> - Dispatcher functions are required to share the exact same positional,
>   optional and keyword-only arguments as their corresponding NumPy
> functions.
>   Otherwise, valid invocations of a NumPy function could result in an
> error when
>   calling its dispatcher.
> - Because default *values* for keyword arguments do not have
>   ``__array_function__`` attributes, by convention we set all default
> argument
>   values to ``None``. This reduces the likelihood of signatures falling out
>   of sync, and minimizes extraneous information in the dispatcher.
>   The only exception should be cases where the argument value in some way
>   effects dispatching, which should be rare.
> An example of the dispatcher for ``np.concatenate`` may be instructive:
> .. code:: python
>     def _concatenate_dispatcher(arrays, axis=None, out=None):
>         for array in arrays:
>             yield array
>         if out is not None:
>             yield out
> The concatenate dispatcher is written as generator function, which allows
> it
> to potentially include the value of the optional ``out`` argument without
> needing to create a new sequence with the (potentially long) list of
> objects
> to be concatenated.
> Trying ``__array_function__`` methods until the right one works
> '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
> Many arguments may implement the ``__array_function__`` protocol. Some
> of these may decide that, given the available inputs, they are unable to
> determine the correct result. How do we call the right one? If several
> are valid then which has precedence?
> For the most part, the rules for dispatch with ``__array_function__``
> match those for ``__array_ufunc__`` (see
> `NEP-13 <https://www.numpy.org/neps/nep-0013-ufunc-overrides.html>`_).
> In particular:
> -  NumPy will gather implementations of ``__array_function__`` from all
>    specified inputs and call them in order: subclasses before
>    superclasses, and otherwise left to right. Note that in some edge cases
>    involving subclasses, this differs slightly from the
>    `current behavior <https://bugs.python.org/issue30140>`_ of Python.
> -  Implementations of ``__array_function__`` indicate that they can
>    handle the operation by returning any value other than
>    ``NotImplemented``.
> -  If all ``__array_function__`` methods return ``NotImplemented``,
>    NumPy will raise ``TypeError``.
> One deviation from the current behavior of ``__array_ufunc__`` is that
> NumPy
> will only call ``__array_function__`` on the *first* argument of each
> unique
> type. This matches Python's
> `rule for calling reflected methods <https://docs.python.org/3/
> reference/datamodel.html#object.__ror__>`_,
> and this ensures that checking overloads has acceptable performance even
> when
> there are a large number of overloaded arguments. To avoid long-term
> divergence
> between these two dispatch protocols, we should
> `also update <https://github.com/numpy/numpy/issues/11306>`_
> ``__array_ufunc__`` to match this behavior.
> Special handling of ``numpy.ndarray``
> '''''''''''''''''''''''''''''''''''''
> The use cases for subclasses with ``__array_function__`` are the same as
> those
> with ``__array_ufunc__``, so ``numpy.ndarray`` should also define a
> ``__array_function__`` method mirroring ``ndarray.__array_ufunc__``:
> .. code:: python
>     def __array_function__(self, func, types, args, kwargs):
>         # Cannot handle items that have __array_function__ other than our
> own.
>         for t in types:
>             if (hasattr(t, '__array_function__') and
>                     t.__array_function__ is not
> ndarray.__array_function__):
>                 return NotImplemented
>         # Arguments contain no overrides, so we can safely call the
>         # overloaded function again.
>         return func(*args, **kwargs)
> To avoid infinite recursion, the dispatch rules for ``__array_function__``
> need
> also the same special case they have for ``__array_ufunc__``: any
> arguments with
> an ``__array_function__`` method that is identical to
> ``numpy.ndarray.__array_function__`` are not be called as
> ``__array_function__`` implementations.
> Changes within NumPy functions
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> Given a function defining the above behavior, for now call it
> ``try_array_function_override``, we now need to call that function from
> within every relevant NumPy function. This is a pervasive change, but of
> fairly simple and innocuous code that should complete quickly and
> without effect if no arguments implement the ``__array_function__``
> protocol.
> In most cases, these functions should written using the
> ``array_function_dispatch`` decorator, which also associates dispatcher
> functions:
> .. code:: python
>     def array_function_dispatch(dispatcher):
>         """Wrap a function for dispatch with the __array_function__
> protocol."""
>         def decorator(func):
>             @functools.wraps(func)
>             def new_func(*args, **kwargs):
>                 relevant_arguments = dispatcher(*args, **kwargs)
>                 success, value = try_array_function_override(
>                     new_func, relevant_arguments, args, kwargs)
>                 if success:
>                     return value
>                 return func(*args, **kwargs)
>             return new_func
>         return decorator
>     # example usage
>     def _broadcast_to_dispatcher(array, shape, subok=None,
> **ignored_kwargs):
>         return (array,)
>     @array_function_dispatch(_broadcast_to_dispatcher)
>     def broadcast_to(array, shape, subok=False):
>         ...  # existing definition of np.broadcast_to
> Using a decorator is great! We don't need to change the definitions of
> existing NumPy functions, and only need to write a few additional lines
> for the dispatcher function. We could even reuse a single dispatcher for
> families of functions with the same signature (e.g., ``sum`` and ``prod``).
> For such functions, the largest change could be adding a few lines to the
> docstring to note which arguments are checked for overloads.
> It's particularly worth calling out the decorator's use of
> ``functools.wraps``:
> - This ensures that the wrapped function has the same name and docstring as
>   the wrapped NumPy function.
> - On Python 3, it also ensures that the decorator function copies the
> original
>   function signature, which is important for introspection based tools
> such as
>   auto-complete. If we care about preserving function signatures on Python
> 2,
>   for the `short while longer <http://www.numpy.org/neps/
> nep-0014-dropping-python2.7-proposal.html>`_
>   that NumPy supports Python 2.7, we do could do so by adding a vendored
>   dependency on the (single-file, BSD licensed)
>   `decorator library <https://github.com/micheles/decorator>`_.
> - Finally, it ensures that the wrapped function
>   `can be pickled <http://gael-varoquaux.info/programming/decoration-in-
> python-done-right-decorating-and-pickling.html>`_.
> In a few cases, it would not make sense to use the
> ``array_function_dispatch``
> decorator directly, but override implementation in terms of
> ``try_array_function_override`` should still be straightforward.
> - Functions written entirely in C (e.g., ``np.concatenate``) can't use
>   decorators, but they could still use a C equivalent of
>   ``try_array_function_override``. If performance is not a concern, they
> could
>   also be easily wrapped with a small Python wrapper.
> - The ``__call__`` method of ``np.vectorize`` can't be decorated with
> <p style="margin:0px;font-stretch:normal;font-size:17.4px;line-
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion at python.org
> https://mail.python.org/mailman/listinfo/numpy-discussion
> I would like to propose that we use `__array_function` in the following
> manner for functions that create arrays:
>    - `array_reference` for indicating the “reference array” whose
>    `__array_function__` implementation will be called. For example,
>    `np.arange(5, array_reference=some_dask_array)`.
>    - I use a reference in the design rather than a type because for some
>    arrays (such as Dask), chunk sizes or other reference data is needed to
>    make this work.
> I realise that this is a big design decision, so I welcome any input!
> Best Regards,
> Hameer Abbasi
> Sent from Astro <https://www.helloastro.com> for Mac
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion at python.org
> https://mail.python.org/mailman/listinfo/numpy-discussion
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20180627/d45cac21/attachment-0001.html>

More information about the NumPy-Discussion mailing list