NEP 35 — 数组创建分派与 __array_function__#

作者:

Peter Andreas Entschev <pentschev@nvidia.com>

状态:

最终

类型:

标准追踪

创建日期:

2019-10-15

更新日期:

2020-11-06

决议:

https://mail.python.org/pipermail/numpy-discussion/2021-May/081761.html

摘要#

我们提议向所有数组创建函数引入一个新的关键字参数 like=，以解决 NEP 18 [1] 中描述的 __array_function__ 的一个缺点。like= 关键字参数将创建参数类型的一个实例，从而能够直接创建非 NumPy 数组。目标数组类型必须实现 __array_function__ 协议。

动机与范围#

许多库都实现了 NumPy API，例如用于图计算的 Dask、用于 GPGPU 计算的 CuPy、用于 N 维标记数组的 xarray 等。在底层，它们都采用了 __array_function__ 协议，这使得 NumPy 能够将下游对象理解并视为原生 numpy.ndarray 对象。因此，社区在使用各种库时仍然受益于统一的 NumPy API。这不仅为标准化带来了极大的便利，还消除了学习新 API 和为每个新对象重写代码的负担。从更专业的角度来说，这种协议机制被称为“分派器”（dispatcher），这是我们从现在开始指代此机制时使用的术语。

x = dask.array.arange(5)    # Creates dask.array
np.diff(x)                  # Returns dask.array

请注意上面我们如何通过调用 np.diff 经由 NumPy 命名空间来调用 Dask 的 diff 实现，同样地，如果我们有一个 CuPy 数组或任何其他采用 __array_function__ 的库中的数组，也会适用。这允许编写与实现库无关的代码，因此用户可以编写一次代码，然后根据需要使用不同的数组实现。

显然，如果数组是在其他地方创建并交由 NumPy 处理，那么有一个就绪的协议是有用的。但是，这些数组仍然必须在其原生库中启动并带回来。如果能够通过 NumPy API 创建这些对象，那么将获得几乎完整的体验，所有操作都使用 NumPy 语法。例如，假设我们有一些 CuPy 数组 cp_arr，并且想要一个具有单位矩阵的类似 CuPy 数组。我们仍然可以这样写：

x = cupy.identity(3)

相反，更好的方法是只使用 NumPy API，现在可以通过以下方式实现：

x = np.identity(3, like=cp_arr)

仿佛魔法般，x 也将是一个 CuPy 数组，因为 NumPy 能够从 cp_arr 的类型中推断出来。请注意，如果没有 like=，最后这一步将无法实现，因为 NumPy 仅凭整数输入将无法知道用户期望的是 CuPy 数组。

提议的新关键字 like= 仅用于识别要分派到的下游库，并且该对象仅用作参考，这意味着不会对该对象执行任何修改、复制或处理。

我们期望此功能对库开发者特别有用，它允许他们根据用户传入的数组创建用于内部用途的新数组，从而避免不必要地创建 NumPy 数组，这些数组最终会导致额外转换为下游数组类型。

NumPy 1.17 之后已放弃对 Python 2.7 的支持，因此我们利用 PEP-3102 [2] 中描述的仅关键字参数标准来实现 like=，从而防止它通过位置参数传递。

用法与影响#

不使用下游库中其他数组的 NumPy 用户可以继续使用不带 like= 参数的数组创建例程。使用 like=np.ndarray 将像没有通过该参数传递数组一样工作。然而，这将导致额外的检查，从而对性能产生负面影响。

为了理解 like= 的预期用途，在进入更复杂的案例之前，请考虑以下仅包含 NumPy 和 CuPy 数组的示例：

import numpy as np
import cupy

def my_pad(arr, padding):
    padding = np.array(padding, like=arr)
    return np.concatenate((padding, arr, padding))

my_pad(np.arange(5), [-1, -1])    # Returns np.ndarray
my_pad(cupy.arange(5), [-1, -1])  # Returns cupy.core.core.ndarray

请注意，在上面的 my_pad 函数中，arr 如何被用作参考，以决定填充应该具有的数组类型，然后才将数组连接以产生结果。另一方面，如果未使用 like=，NumPy 的情况仍然会工作，但 CuPy 不会允许这种自动转换，最终会引发 TypeError: Only cupy arrays can be concatenated 异常。

现在我们应该看看像 Dask 这样的库如何从 like= 中受益。在理解这一点之前，重要的是要了解一些 Dask 基础知识以及它如何确保 __array_function__ 的正确性。请注意，Dask 可以对不同类型的对象执行计算，例如数据帧（dataframes）、包（bags）和数组（arrays），这里我们将严格关注数组，这些是我们可以使用 __array_function__ 的对象。

Dask 使用图计算模型，这意味着它将一个大问题分解为许多小问题，然后合并它们的结果以达到最终结果。为了将问题分解为更小的部分，Dask 还会将数组分解为它称为“块”（chunks）的更小数组。因此，一个 Dask 数组可以由一个或多个块组成，并且它们可以是不同的类型。然而，在 __array_function__ 的上下文中，Dask 只允许相同类型的块；例如，一个 Dask 数组可以由多个 NumPy 数组或多个 CuPy 数组组成，但不能是两者的混合。

为了避免计算过程中类型不匹配，Dask 在整个计算过程中将其数组的一部分保持一个属性 _meta：此属性既用于在图创建时预测输出类型，又用于创建某个函数计算中所需的任何中间数组。回到我们之前的例子，我们可以使用 _meta 信息来识别我们将用于填充的数组类型，如下所示：

import numpy as np
import cupy
import dask.array as da
from dask.array.utils import meta_from_array

def my_dask_pad(arr, padding):
    padding = np.array(padding, like=meta_from_array(arr))
    return np.concatenate((padding, arr, padding))

# Returns dask.array<concatenate, shape=(9,), dtype=int64, chunksize=(5,), chunktype=numpy.ndarray>
my_dask_pad(da.arange(5), [-1, -1])

# Returns dask.array<concatenate, shape=(9,), dtype=int64, chunksize=(5,), chunktype=cupy.ndarray>
my_dask_pad(da.from_array(cupy.arange(5)), [-1, -1])

注意上面返回值中的 chunktype 如何从第一个 my_dask_pad 调用中的 numpy.ndarray 变为第二个调用中的 cupy.ndarray。我们还在本示例中将函数重命名为 my_dask_pad，旨在明确这是 Dask 在需要时将如何实现此类功能的方式，因为它需要 Dask 的内部工具，而这些工具在其他地方用处不大。

为了正确识别数组类型，我们使用 Dask 的实用函数 meta_from_array，该函数作为支持 __array_function__ 的工作的一部分引入，允许 Dask 适当地处理 _meta。读者可以将 meta_from_array 视为一个特殊函数，它只返回底层 Dask 数组的类型，例如：

np_arr = da.arange(5)
cp_arr = da.from_array(cupy.arange(5))

meta_from_array(np_arr)  # Returns a numpy.ndarray
meta_from_array(cp_arr)  # Returns a cupy.ndarray

由于 meta_from_array 返回的值是类 NumPy 数组，我们可以直接将其传递给 like= 参数。

__meta_from_array__ 函数主要用于库的内部使用，以确保创建的块具有正确的类型。如果没有 like= 参数，将无法确保 my_pad 创建的填充数组类型与输入数组的类型匹配，这将导致 CuPy 引发 TypeError 异常，正如前面讨论的，这只会在 CuPy 情况下发生。结合 Dask 对元数组的内部处理和提议的 like= 参数，现在可以处理涉及创建非 NumPy 数组的情况，这可能是 Dask 目前从 __array_function__ 协议中面临的最大限制。

向后兼容性#

本提案不会在 NumPy 中引起任何向后兼容性问题，因为它只是向现有数组创建函数引入了一个默认值为 None 的新关键字参数，因此不会改变当前行为。

详细描述#

__array_function__ 协议的引入使得下游库开发者能够使用 NumPy 作为分派 API。然而，该协议并未（也不打算）解决下游库创建数组的问题，从而阻止这些库在该上下文中利用如此重要的功能。

本 NEP 的目的是以一种简单直接的方式解决这一缺陷：引入一个新的 like= 关键字参数，类似于 empty_like 系列函数的工作方式。当数组创建函数收到此类参数时，它们将触发 __array_function__ 协议，并调用下游库自己的数组创建函数实现。like= 参数，顾名思义，应仅用于识别分派目标。与 __array_function__ 迄今为止的使用方式（第一个参数识别目标下游库）相反，为了避免破坏 NumPy 在数组创建方面的 API，新的 like= 关键字将用于分派目的。

下游库将受益于 like= 参数，而无需对其 API 进行任何更改，因为该参数只需由 NumPy 实现。仍然允许下游库包含 like= 参数，因为它在某些情况下可能有用，有关这些情况的详细信息，请参阅实现。仍将要求下游库实现 NEP 18 [1] 中描述的 __array_function__ 协议，并适当地将其参数引入到其对 NumPy 数组创建函数的调用中，如用法与影响中所示。

实现#

该实现需要在 NumPy 的所有现有数组创建函数中引入一个新的 like= 关键字。可以引用添加此新参数的函数示例（但不限于）包括接受类数组对象（如 array 和 asarray）的函数、基于数值输入创建数组（如 range 和 identity）的函数，以及 empty 系列函数，尽管这可能有些冗余，因为这些函数已经存在带有 empty_like 命名格式的专用版本。截至本 NEP 撰写时，完整的数组创建函数列表可在 [5] 中找到。

这个新提出的关键字在分派之前应由 __array_function__ 机制从关键字字典中移除。这样做的目的是双重的：

  1. 简化已选择实现 __array_function__ 协议的库采用数组创建，从而消除了对所有数组创建函数明确选择加入的要求；以及

  2. 大多数下游库将不需要该关键字参数，而需要该参数的库可以通过从 __array_function__ 中捕获 self 来实现。

因此，下游库不需要在其数组创建 API 中包含 like= 关键字。在某些情况下（例如 Dask），拥有 like= 关键字可能很有用，因为它将允许实现识别数组内部。例如，Dask 可以利用引用数组来识别其块类型（例如 NumPy、CuPy、Sparse），从而创建由相同块类型支持的新 Dask 数组，除非 Dask 可以读取引用数组的属性，否则这是不可能的。

函数分派#

有两种不同的分派情况：Python 函数和 C 函数。为了允许 __array_function__ 分派，一种可能的实现是用 overrides.array_function_dispatch 装饰 Python 函数，但 C 函数有不同的要求，我们将在稍后描述。

下面的例子展示了如何用 overrides.array_function_dispatch 装饰 asarray 的建议：

def _asarray_decorator(a, dtype=None, order=None, *, like=None):
    return (like,)

@set_module('numpy')
@array_function_dispatch(_asarray_decorator)
def asarray(a, dtype=None, order=None, *, like=None):
    return array(a, dtype, copy=False, order=order)

请注意，在上面的示例中，实现保持不变，唯一的区别是装饰器，它使用新的 _asarray_decorator 函数来指示 __array_function__ 协议在 like 不为 None 时进行分派。

现在我们来看一个 C 函数的例子，既然 asarray 无论如何都是 array 的一个特化，我们现在将以后者为例。由于 array 是一个 C 函数，目前 NumPy 对其 Python 源代码所做的只是导入该函数并将其 __module__ 调整为 numpy。该函数现在将用 overrides.array_function_from_dispatcher 的一个特化进行装饰，该特化也将负责调整模块。

array_function_nodocs_from_c_func_and_dispatcher = functools.partial(
    overrides.array_function_from_dispatcher,
    module='numpy', docs_from_dispatcher=False, verify=False)

@array_function_nodocs_from_c_func_and_dispatcher(_multiarray_umath.array)
def array(a, dtype=None, *, copy=True, order='K', subok=False, ndmin=0,
          like=None):
    return (like,)

上述 C 函数实现存在两个缺点：

  1. 它会创建另一个 Python 函数调用；并且

  2. 为了遵循当前的实现标准，文档应直接附加到 Python 源代码中。

本提案的第一个版本建议将上述实现作为 NumPy 中以 C 语言实现的函数的一种可行解决方案。然而，由于上面指出的缺点，我们决定放弃对 Python 端的任何更改，并使用纯 C 实现来解决这些问题。详情请参阅 [7]。

下游读取引用数组#

正如实现部分开头所述，like= 不会传播到下游库，尽管如此，仍然可以访问它。这需要对下游库的 __array_function__ 定义进行一些更改，其中 self 属性实际上是通过 like= 传递的。之所以如此，是因为我们使用 like= 作为分派数组，这与 NEP-18 涵盖的其他计算函数不同，后者通常在第一个位置参数上进行分派。

这种用法的示例如下：在保留后端类型的同时创建一个新的 Dask 数组：

# Returns dask.array<array, shape=(3,), dtype=int64, chunksize=(3,), chunktype=cupy.ndarray>
np.asarray([1, 2, 3], like=da.array(cp.array(())))

# Returns a cupy.ndarray
type(np.asarray([1, 2, 3], like=da.array(cp.array(()))).compute())

注意上面数组是如何由 chunktype=cupy.ndarray 支持的，并且计算后的结果数组也是一个 cupy.ndarray。如果 Dask 没有通过 __array_function__ 的 self 属性使用 like= 参数，上面的例子将改为由 numpy.ndarray 支持。

# Returns dask.array<array, shape=(3,), dtype=int64, chunksize=(3,), chunktype=numpy.ndarray>
np.asarray([1, 2, 3], like=da.array(cp.array(())))

# Returns a numpy.ndarray
type(np.asarray([1, 2, 3], like=da.array(cp.array(()))).compute())

鉴于库需要依赖 __array_function__ 中的 self 属性来分派带正确引用数组的函数，我们建议以下两种替代方案之一：

  1. 在下游库中引入一个支持 like= 参数的函数列表，并在调用函数时传递 like=self；或者

  2. 检查函数的签名并验证它是否包含 like= 参数。请注意，这可能会导致更高的性能开销，并且假设内省是可能的，如果函数是 C 函数，则可能无法进行内省。

为了使事情更清晰，让我们看看建议 2 如何在 Dask 中实现。__array_function__ 定义中与 Dask 相关的当前部分如下所示：

def __array_function__(self, func, types, args, kwargs):
    # Code not relevant for this example here

    # Dispatch ``da_func`` (da.asarray, for example) with *args and **kwargs
    da_func(*args, **kwargs)

更新后的代码如下所示：

def __array_function__(self, func, types, args, kwargs):
    # Code not relevant for this example here

    # Inspect ``da_func``'s  signature and store keyword-only arguments
    import inspect
    kwonlyargs = inspect.getfullargspec(da_func).kwonlyargs

    # If ``like`` is contained in ``da_func``'s signature, add ``like=self``
    # to the kwargs dictionary.
    if 'like' in kwonlyargs:
        kwargs['like'] = self

    # Dispatch ``da_func`` (da.asarray, for example) with args and kwargs.
    # Here, kwargs contain ``like=self`` if the function's signature does too.
    da_func(*args, **kwargs)

替代方案#

最近，NEP 37 [6] 提出了一个完全取代 __array_function__ 的新协议，这将需要已采用 __array_function__ 的下游库进行大量返工，因此我们仍然认为 like= 参数对 NumPy 和下游库有利。然而，该提案不一定被视为本 NEP 的直接替代方案，因为它将完全取代 NEP 18，而本 NEP 正是基于 NEP 18 构建的。关于此新提案的详细信息以及为何需要下游库进行返工的讨论超出了本提案的范围。

讨论#

参考文献#