Merge pull request #13979 from danielballan/array-function-high-level-docs

DOC: Document array_function at a higher level.

4 files changed, 404 insertions, 2 deletions

diff --git a/doc/source/reference/arrays.classes.rst b/doc/source/reference/arrays.classes.rst
index 3b13530c7..a91215476 100644
--- a/doc/source/reference/arrays.classes.rst
+++ b/doc/source/reference/arrays.classes.rst
@@ -6,8 +6,15 @@ Standard array subclasses
 
 .. currentmodule:: numpy
 
-The :class:`ndarray` in NumPy is a "new-style" Python
-built-in-type. Therefore, it can be inherited from (in Python or in C)
+.. note::
+
+    Subclassing a ``numpy.ndarray`` is possible but if your goal is to create
+    an array with *modified* behavior, as do dask arrays for distributed
+    computation and cupy arrays for GPU-based computation, subclassing is
+    discouraged. Instead, using numpy's
+    :ref:`dispatch mechanism <basics.dispatch>` is recommended.
+
+The :class:`ndarray` can be inherited from (in Python or in C)
 if desired. Therefore, it can form a foundation for many useful
 classes. Often whether to sub-class the array object or to simply use
 the core array component as an internal part of a new class is a
@@ -147,6 +154,121 @@ NumPy provides several hooks that classes can customize:
       :func:`__array_prepare__`, :data:`__array_priority__` mechanism
       described below for ufuncs (which may eventually be deprecated).
 
+.. py:method:: class.__array_function__(func, types, args, kwargs)
+
+   .. versionadded:: 1.16
+
+   .. note::
+
+       - In NumPy 1.17, the protocol is enabled by default, but can be disabled
+         with ``NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=0``.
+       - In NumPy 1.16, you need to set the environment variable
+         ``NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=1`` before importing NumPy to use
+         NumPy function overrides.
+       - Eventually, expect to ``__array_function__`` to always be enabled.
+
+   -  ``func`` is an arbitrary callable exposed by NumPy's public API,
+      which was called in the form ``func(*args, **kwargs)``.
+   -  ``types`` is a `collection <collections.abc.Collection>`_
+      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.
+
+   As a convenience for ``__array_function__`` implementors, ``types``
+   provides all argument types with an ``'__array_function__'`` attribute.
+   This allows implementors to quickly identify cases where they should defer
+   to ``__array_function__`` implementations on other arguments.
+   Implementations should not rely on the iteration order of ``types``.
+
+   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.
+
+   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.
+
+   It may also be convenient to define a custom decorators (``implements``
+   below) for registering ``__array_function__`` implementations.
+
+   .. code:: python
+
+       HANDLED_FUNCTIONS = {}
+
+       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.
+
+   Just like the case for builtin special methods like ``__add__``, properly
+   written ``__array_function__`` methods should always return
+   ``NotImplemented`` when an unknown type is encountered. Otherwise, it will
+   be impossible to correctly override NumPy functions from another object
+   if the operation also includes one of your objects.
+
+   For the most part, the rules for dispatch with ``__array_function__``
+   match those for ``__array_ufunc__``. 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``.
+
+   If no ``__array_function__`` methods exists, NumPy will default to calling
+   its own implementation, intended for use on NumPy arrays. This case arises,
+   for example, when all array-like arguments are Python numbers or lists.
+   (NumPy arrays do have a ``__array_function__`` method, given below, but it
+   always returns ``NotImplemented`` if any argument other than a NumPy array
+   subclass implements ``__array_function__``.)
+
+   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.
+
 .. py:method:: class.__array_finalize__(obj)
 
    This method is called whenever the system internally allocates a
diff --git a/doc/source/user/basics.dispatch.rst b/doc/source/user/basics.dispatch.rst
new file mode 100644
index 000000000..f7b8da262
--- /dev/null
+++ b/doc/source/user/basics.dispatch.rst
@@ -0,0 +1,8 @@
+.. _basics.dispatch:
+
+*******************************
+Writing custom array containers
+*******************************
+
+.. automodule:: numpy.doc.dispatch
+
diff --git a/doc/source/user/basics.rst b/doc/source/user/basics.rst
index 7875aff6e..e0fc0ece3 100644
--- a/doc/source/user/basics.rst
+++ b/doc/source/user/basics.rst
@@ -12,4 +12,5 @@ NumPy basics
    basics.broadcasting
    basics.byteswapping
    basics.rec
+   basics.dispatch
    basics.subclassing
diff --git a/numpy/doc/dispatch.py b/numpy/doc/dispatch.py
new file mode 100644
index 000000000..09a3e5134
--- /dev/null
+++ b/numpy/doc/dispatch.py
@@ -0,0 +1,271 @@
+""".. _dispatch_mechanism:
+
+Numpy's dispatch mechanism, introduced in numpy version v1.16 is the
+recommended approach for writing custom N-dimensional array containers that are
+compatible with the numpy API and provide custom implementations of numpy
+functionality. Applications include `dask <http://dask.pydata.org>`_ arrays, an
+N-dimensional array distributed across multiple nodes, and `cupy
+<https://docs-cupy.chainer.org/en/stable/>`_ arrays, an N-dimensional array on
+a GPU.
+
+To get a feel for writing custom array containers, we'll begin with a simple
+example that has rather narrow utility but illustrates the concepts involved.
+
+>>> import numpy as np
+>>> class DiagonalArray:
+...     def __init__(self, N, value):
+...         self._N = N
+...         self._i = value
+...     def __repr__(self):
+...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
+...     def __array__(self):
+...         return self._i * np.eye(self._N)
+...
+
+Our custom array can be instantiated like:
+
+>>> arr = DiagonalArray(5, 1)
+>>> arr
+DiagonalArray(N=5, value=1)
+
+We can convert to a numpy array using :func:`numpy.array` or
+:func:`numpy.asarray`, which will call its ``__array__`` method to obtain a
+standard ``numpy.ndarray``.
+
+>>> np.asarray(arr)
+array([[1., 0., 0., 0., 0.],
+       [0., 1., 0., 0., 0.],
+       [0., 0., 1., 0., 0.],
+       [0., 0., 0., 1., 0.],
+       [0., 0., 0., 0., 1.]])
+
+If we operate on ``arr`` with a numpy function, numpy will again use the
+``__array__`` interface to convert it to an array and then apply the function
+in the usual way.
+
+>>> np.multiply(arr, 2)
+array([[2., 0., 0., 0., 0.],
+       [0., 2., 0., 0., 0.],
+       [0., 0., 2., 0., 0.],
+       [0., 0., 0., 2., 0.],
+       [0., 0., 0., 0., 2.]])
+
+
+Notice that the return type is a standard ``numpy.ndarray``.
+
+>>> type(arr)
+numpy.ndarray
+
+How can we pass our custom array type through this function? Numpy allows a
+class to indicate that it would like to handle computations in a custom-defined
+way through the interaces ``__array_ufunc__`` and ``__array_function__``. Let's
+take one at a time, starting with ``_array_ufunc__``. This method covers
+:ref:`ufuncs`, a class of functions that includes, for example,
+:func:`numpy.multiply` and :func:`numpy.sin`.
+
+The ``__array_ufunc__`` receives:
+
+- ``ufunc``, a function like ``numpy.multiply``
+- ``method``, a string, differentiating between ``numpy.multiply(...)`` and
+  variants like ``numpy.multiply.outer``, ``numpy.multiply.accumulate``, and so
+  on.  For the common case, ``numpy.multiply(...)``, ``method == '__call__'``.
+- ``inputs``, which could be a mixture of different types
+- ``kwargs``, keyword arguments passed to the function
+
+For this example we will only handle the method ``'__call__``.
+
+>>> from numbers import Number
+>>> class DiagonalArray:
+...     def __init__(self, N, value):
+...         self._N = N
+...         self._i = value
+...     def __repr__(self):
+...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
+...     def __array__(self):
+...         return self._i * np.eye(self._N)
+...     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+...         if method == '__call__':
+...             N = None
+...             scalars = []
+...             for input in inputs:
+...                 if isinstance(input, Number):
+...                     scalars.append(input)
+...                 elif isinstance(input, self.__class__):
+...                     scalars.append(input._i)
+...                     if N is not None:
+...                         if N != self._N:
+...                             raise TypeError("inconsistent sizes")
+...                     else:
+...                         N = self._N
+...                 else:
+...                     return NotImplemented
+...             return self.__class__(N, ufunc(*scalars, **kwargs))
+...         else:
+...             return NotImplemented
+...
+
+Now our custom array type passes through numpy functions.
+
+>>> arr = DiagonalArray(5, 1)
+>>> np.multiply(arr, 3)
+DiagonalArray(N=5, value=3)
+>>> np.add(arr, 3)
+DiagonalArray(N=5, value=4)
+>>> np.sin(arr)
+DiagonalArray(N=5, value=0.8414709848078965)
+
+At this point ``arr + 3`` does not work.
+
+>>> arr + 3
+TypeError: unsupported operand type(s) for *: 'DiagonalArray' and 'int'
+
+To support it, we need to define the Python interfaces ``__add__``, ``__lt__``,
+and so on to dispatch to the corresponding ufunc. We can achieve this
+conveniently by inheriting from the mixin
+:class:`~numpy.lib.mixins.NDArrayOperatorsMixin`.
+
+>>> import numpy.lib.mixins
+>>> class DiagonalArray(numpy.lib.mixins.NDArrayOperatorsMixin):
+...     def __init__(self, N, value):
+...         self._N = N
+...         self._i = value
+...     def __repr__(self):
+...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
+...     def __array__(self):
+...         return self._i * np.eye(self._N)
+...     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+...         if method == '__call__':
+...             N = None
+...             scalars = []
+...             for input in inputs:
+...                 if isinstance(input, Number):
+...                     scalars.append(input)
+...                 elif isinstance(input, self.__class__):
+...                     scalars.append(input._i)
+...                     if N is not None:
+...                         if N != self._N:
+...                             raise TypeError("inconsistent sizes")
+...                     else:
+...                         N = self._N
+...                 else:
+...                     return NotImplemented
+...             return self.__class__(N, ufunc(*scalars, **kwargs))
+...         else:
+...             return NotImplemented
+...
+
+>>> arr = DiagonalArray(5, 1)
+>>> arr + 3
+DiagonalArray(N=5, value=4)
+>>> arr > 0
+DiagonalArray(N=5, value=True)
+
+Now let's tackle ``__array_function__``. We'll create dict that maps numpy
+functions to our custom variants.
+
+>>> HANDLED_FUNCTIONS = {}
+>>> class DiagonalArray(numpy.lib.mixins.NDArrayOperatorsMixin):
+...     def __init__(self, N, value):
+...         self._N = N
+...         self._i = value
+...     def __repr__(self):
+...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
+...     def __array__(self):
+...         return self._i * np.eye(self._N)
+...     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+...         if method == '__call__':
+...             N = None
+...             scalars = []
+...             for input in inputs:
+...                 # In this case we accept only scalar numbers or DiagonalArrays.
+...                 if isinstance(input, Number):
+...                     scalars.append(input)
+...                 elif isinstance(input, self.__class__):
+...                     scalars.append(input._i)
+...                     if N is not None:
+...                         if N != self._N:
+...                             raise TypeError("inconsistent sizes")
+...                     else:
+...                         N = self._N
+...                 else:
+...                     return NotImplemented
+...             return self.__class__(N, ufunc(*scalars, **kwargs))
+...         else:
+...             return NotImplemented
+...    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 DiagonalArray objects.
+...        if not all(issubclass(t, self.__class__) for t in types):
+...            return NotImplemented
+...        return HANDLED_FUNCTIONS[func](*args, **kwargs)
+...
+
+A convenient pattern is to define a decorator ``implements`` that can be used
+to add functions to ``HANDLED_FUNCTIONS``.
+
+>>> def implements(np_function):
+...    "Register an __array_function__ implementation for DiagonalArray objects."
+...    def decorator(func):
+...        HANDLED_FUNCTIONS[np_function] = func
+...        return func
+...    return decorator
+...
+
+Now we write implementations of numpy functions for ``DiagonalArray``.
+For completeness, to support the usage ``arr.sum()`` add a method ``sum`` that
+calls ``numpy.sum(self)``, and the same for ``mean``.
+
+>>> @implements(np.sum)
+... def sum(a):
+...     "Implementation of np.sum for DiagonalArray objects"
+...     return arr._i * arr._N
+...
+>>> @implements(np.mean)
+... def sum(a):
+...     "Implementation of np.mean for DiagonalArray objects"
+...     return arr._i / arr._N
+...
+>>> arr = DiagonalArray(5, 1)
+>>> np.sum(arr)
+5
+>>> np.mean(arr)
+0.2
+
+If the user tries to use any numpy functions not included in
+``HANDLED_FUNCTIONS``, a ``TypeError`` will be raised by numpy, indicating that
+this operation is not supported. For example, concatenating two
+``DiagonalArrays`` does not produce another diagonal array, so it is not
+supported.
+
+>>> np.concatenate([arr, arr])
+TypeError: no implementation found for 'numpy.concatenate' on types that implement __array_function__: [<class '__main__.DiagonalArray'>]
+
+Additionally, our implementations of ``sum`` and ``mean`` do not accept the
+optional arguments that numpy's implementation does.
+
+>>> np.sum(arr, axis=0)
+TypeError: sum() got an unexpected keyword argument 'axis'
+
+The user always has the option of converting to a normal ``numpy.ndarray`` with
+:func:`numpy.asarray` and using standard numpy from there.
+
+>>> np.concatenate([np.asarray(arr), np.asarray(arr)])
+array([[1., 0., 0., 0., 0.],
+       [0., 1., 0., 0., 0.],
+       [0., 0., 1., 0., 0.],
+       [0., 0., 0., 1., 0.],
+       [0., 0., 0., 0., 1.],
+       [1., 0., 0., 0., 0.],
+       [0., 1., 0., 0., 0.],
+       [0., 0., 1., 0., 0.],
+       [0., 0., 0., 1., 0.],
+       [0., 0., 0., 0., 1.]])
+
+Refer to the `dask source code <https://github.com/dask/dask>`_ and
+`cupy source code <https://github.com/cupy/cupy>`_  for more fully-worked
+examples of custom array containers.
+
+See also `NEP 18 <http://www.numpy.org/neps/nep-0018-array-function-protocol.html>`_.
+"""