Interop: array protocols and the array API

NumPy · Expert · 18 min

Modern code runs on NumPy, CuPy, Dask, JAX or PyTorch; these protocols decide which library actually does the work.

NumPy offers several hooks, from least to most powerful, through which other types interact with it:

  • `__array__`: lets np.asarray(obj) convert the object to an ndarray. In NumPy 2 it is called with dtype= and copy= keywords, and implementations should accept copy so that copy=False can be honoured or refused.
  • `__array_ufunc__` (NEP 13): when any operand defines it, calling a ufunc such as np.add(x, y) or np.sin(x) is handed to that method instead of converting to ndarray. Dask, CuPy, xarray and unit libraries like Pint use it to keep their own types (lazy graphs, GPU arrays, quantities).
  • `__array_function__` (NEP 18): the same idea for non-ufunc functions like np.concatenate or np.mean. The implementation receives the NumPy function and dispatches to its own version, or returns NotImplemented.

These protocols made "write NumPy, run elsewhere" partly possible, but they leak: functions that create arrays from scratch (like np.zeros) have no array argument to dispatch on, and coverage differs per library.

The more robust modern approach is the Python array API standard, a common specification of functions and semantics that NumPy 2 implements in its main namespace (x.__array_namespace__() returns the module). Library code gets the right namespace from its inputs and uses only standard functions:

xp = array_namespace(x)   # from the array_api_compat package
y = xp.exp(x) / xp.sum(xp.exp(x))

The same code then runs unchanged on NumPy, CuPy, PyTorch or JAX arrays. SciPy and scikit-learn are adopting this pattern. Stick to the standard's subset (no in-place view tricks, no np.matrix, explicit dtypes) when portability matters.

Example

import numpy as np

class Logged:  # wraps an ndarray, reports which NumPy calls hit it
    def __init__(self, data):
        self.data = np.asarray(data)
    def __array__(self, dtype=None, copy=None):
        return self.data if dtype is None else self.data.astype(dtype)
    def __array_ufunc__(self, ufunc, method, *inputs, **kw):
        args = [i.data if isinstance(i, Logged) else i for i in inputs]
        print("ufunc:", ufunc.__name__, method)
        return Logged(getattr(ufunc, method)(*args, **kw))
    def __array_function__(self, func, types, args, kwargs):
        print("function:", func.__name__)
        args = [a.data if isinstance(a, Logged) else a for a in args]
        return func(*args, **kwargs)

x = Logged([1.0, 2.0, 3.0])
y = np.sqrt(x)              # ufunc: sqrt __call__
print(np.mean(y))           # function: mean, then 1.382...
print(np.asarray(x) * 2)    # [2. 4. 6.]

a = np.ones(3)
xp = a.__array_namespace__()  # NumPy's array-API namespace
print(xp.sum(a))              # 3.0

Remember

  • __array__ converts; __array_ufunc__ and __array_function__ take over NumPy calls
  • Dask, CuPy, xarray and Pint rely on these overrides
  • Array-creation functions can't dispatch on inputs; the array API standard avoids that gap
  • Get xp from your inputs and use standard functions for library-agnostic code

Watch out for

  • Implementing __array__ without the copy keyword, which triggers a DeprecationWarning in NumPy 2 when callers pass copy=.

Check yourself

  1. A Dask array `d` is passed to `np.sin(d)`. Why does the result stay a lazy Dask array?
  2. Why is `np.zeros(n)` a problem for __array_function__-based dispatch?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: numpy.org/doc/stable/user/basics.interoperability.html