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= 关键字应用于分派目的。

下游库无需更改其 API 即可受益于 like= 参数，因为该参数只需要由 NumPy 实现。下游库仍然允许包含 like= 参数，因为它在某些情况下很有用，请参阅 实现 以了解这些情况的详细信息。下游库仍然需要实现 __array_function__ 协议，如 NEP 18 [1] 中所述，并在其对 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 函数，则可能不会。

为了使事情更清楚，让我们看看如何在 Dask 中实现建议 2。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 18，而本 NEP 则在此基础上构建。关于此新提案的详细信息以及为什么这将需要下游库进行返工的讨论超出了本提案的范围。

讨论#

参考文献#

本页内容