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 和重写代码的负担。更技术性的说法是，该协议的这种机制称为“调度程序”，这是我们从现在开始在提及该机制时使用的术语。

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

请注意上面我们如何通过 NumPy 命名空间调用 Dask 的 diff 实现，方法是调用 np.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 可以对不同类型的对象执行计算，例如数据帧、包和数组，在这里我们将严格关注数组，这是我们可以使用 __array_function__ 的对象。

Dask 使用图计算模型，这意味着它将一个大问题分解成许多较小的问题，并将它们的结果合并起来以达到最终结果。为了将问题分解成更小的问题，Dask 还将数组分解成它称为“块”的更小数组。因此，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__ 协议中面临的最大限制。

向后兼容性#

鉴于此提案仅为现有的数组创建函数引入一个新的关键字参数，其默认值为 None，因此不会改变当前行为，所以此提案不会引发 NumPy 中的任何向后兼容性问题。

详细描述#

__array_function__ 协议的引入允许下游库开发人员使用 NumPy 作为调度 API。但是，该协议没有（并且不打算）解决下游库创建数组的问题，从而阻止这些库在该上下文中使用此类重要功能。

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

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

实现#

实现需要为 NumPy 的所有现有数组创建函数引入一个新的 like= 关键字。作为将添加此新参数的函数示例（但不限于此），我们可以引用那些接受类似数组的对象（例如 array 和 asarray），基于数值输入创建数组的函数（例如 range 和 identity），以及 empty 函数系列，即使这可能是多余的，因为这些函数已经存在具有 empty_like 命名格式的专用函数。在本 NEP 撰写之时，可以在 [5] 中找到完整的数组创建函数列表。

这个新提出的关键字将在调度之前由 __array_function__ 机制从关键字字典中删除。这样做的目的是双重的

简化那些已经选择实现 __array_function__ 协议的库对数组创建的采用，从而无需显式选择加入所有数组创建函数；以及
大多数下游库不需要关键字参数，而那些需要的库可以通过从 __array_function__ 中捕获 self 来实现。

因此，下游库不需要将其 like= 关键字添加到其数组创建 API 中。在某些情况下（例如 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 函数的实现有两个缺点

它创建了另一个 Python 函数调用；以及
为了遵循当前的实现标准，文档应直接附加到 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属性来使用正确的参考数组调度函数，我们建议以下两种替代方案之一：

在下游库中引入一个支持like=参数的函数列表，并在调用函数时传递like=self；或者
检查函数的签名，并验证其是否包含like=参数。请注意，这可能会导致更高的性能损失，并假设可以进行内省，如果函数是C函数，则可能无法进行内省。

为了更清楚地说明，让我们看看建议2如何在Dask中实现。Dask中__array_function__定义的当前相关部分如下所示：

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所基于的NEP 18。关于此新提案的详细信息以及为什么这需要下游库进行返工的讨论超出了本提案的范围。

讨论#

参考文献#

版权#

本文档已进入公有领域。