NEP 31 — NumPy API 的上下文局部和全局覆盖

作者:

Hameer Abbasi <habbasi@quansight.com>

作者:

Ralf Gommers <rgommers@quansight.com>

作者:

Peter Bell <pbell@quansight.com>

状态:

已废弃

已替换:

NEP 56 — NumPy 主命名空间中的数组 API 标准支持

类型:

标准轨道

创建时间:

2019-08-22

解决时间:

https://mail.python.org/archives/list/numpy-discussion@python.org/message/Z6AA5CL47NHBNEPTFWYOTSUVSRDGHYPN/

摘要

此 NEP 提议通过一个可扩展的后端机制使 NumPy 的所有公共 API 都可以被覆盖。

接受此 NEP 意味着 NumPy 将在一个独立的命名空间中提供全局和上下文局部的覆盖，以及一个类似于 NEP-18 的分发机制 [2]。__array_function__ 的初步经验表明，有必要能够覆盖那些不接受类数组参数的 NumPy 函数，因此它们无法通过 __array_function__ 来覆盖。最迫切的需求是数组创建和类型转换函数，例如 numpy.zeros 或 numpy.asarray；请参见 NEP-30 [9]。

此 NEP 提议以一种选择加入（opt-in）的方式，允许覆盖 NumPy API 的任何部分。它旨在全面解决 NEP-22 [3]，并消除了为每种需要可覆盖的新函数或对象类型添加一个不断增长的新协议列表的需要。

动机与范围

此 NEP 的主要最终目标是实现以下功能：

# On the library side
import numpy.overridable as unp

def library_function(array):
    array = unp.asarray(array)
    # Code using unumpy as usual
    return array

# On the user side:
import numpy.overridable as unp
import uarray as ua
import dask.array as da

ua.register_backend(da) # Can be done within Dask itself

library_function(dask_array)  # works and returns dask_array

with unp.set_backend(da):
    library_function([1, 2, 3, 4])  # actually returns a Dask array.

在此，backend 可以是 NumPy 或外部库（如 Dask 或 CuPy）定义的任何兼容对象。理想情况下，它应该是 dask.array 或 cupy 模块本身。

这些类型的覆盖对于最终用户和库作者都很有用。最终用户可能已经编写或希望编写代码，然后他们可以加速或迁移到不同的实现，例如 PyData/Sparse。他们可以通过设置一个后端来实现这一点。库作者也可能希望编写可在不同数组实现之间移植的代码，例如 sklearn 可能希望为机器学习算法编写可在不同数组实现之间移植的代码，同时还使用数组创建函数。

此 NEP 采取了一种整体方法：它假定 API 中存在需要可覆盖的部分，并且这些部分会随着时间的推移而增长。它提供了一个通用框架和一个机制，以避免每次需要时都设计新的协议。这是 uarray 的目标：允许在 API 中进行覆盖，而无需设计新的协议。

此 NEP 提议：将 unumpy [8] 作为尚未被 __array_function__ 或 __array_ufunc__ 覆盖的 NumPy API 部分的推荐覆盖机制，并将 uarray 迁入 NumPy 的新命名空间中，以便用户和下游依赖项可以访问这些覆盖。这种迁入机制类似于 SciPy 为使 scipy.fft 可覆盖所做的决定（参见 [10]）。

uarray 背后的动机是多方面的：首先，有几种尝试允许分发 NumPy API 的部分，包括（最突出的是）NEP-13 中的 __array_ufunc__ 协议 [4]，以及 NEP-18 中的 __array_function__ 协议 [2]，但事实证明这需要开发更多协议，包括类型转换协议（参见 [5], [9]）。在参考文献中已经广泛讨论了需要这些覆盖的原因，此 NEP 不会试图深入探讨为什么需要这些；但简而言之：库作者必须能够将任意对象转换为他们自己的类型的数组，例如 CuPy 需要转换为 CuPy 数组，而不是 NumPy 数组。简而言之，需要像 np.asarray(...) 这样的功能，或者“即插即用”并返回鸭子数组（duck-arrays）的替代方案。

用法与影响

此 NEP 允许全局和上下文局部的覆盖，以及类似 __array_function__ 的自动覆盖。

除了在动机部分所述的第一个之外，此 NEP 还将实现以下一些用例：

第一个是允许不同的数据类型（dtypes）返回它们各自的数组。

# Returns an XND array
x = unp.ones((5, 5), dtype=xnd_dtype) # Or torch dtype

第二个是允许覆盖 API 的部分。这是为了允许 np.linalg、BLAS 和 np.random 的替代和/或优化实现。

import numpy as np
import pyfftw # Or mkl_fft

# Makes pyfftw the default for FFT
np.set_global_backend(pyfftw)

# Uses pyfftw without monkeypatching
np.fft.fft(numpy_array)

with np.set_backend(pyfftw) # Or mkl_fft, or numpy
    # Uses the backend you specified
    np.fft.fft(numpy_array)

这将允许一种官方方式来实现覆盖，而无需 monkeypatching 或分发 NumPy 的修改版本。

这里有一些其他用例，虽然没有明确说明但隐含在内：

data = da.from_zarr('myfile.zarr')
# result should still be dask, all things being equal
result = library_function(data)
result.to_zarr('output.zarr')

如果 magic_library 是建立在 unumpy 之上的，那么第二个例子将会起作用。

from dask import array as da
from magic_library import pytorch_predict

data = da.from_zarr('myfile.zarr')
# normally here one would use e.g. data.map_overlap
result = pytorch_predict(data)
result.to_zarr('output.zarr')

有些后端可能依赖于其他后端，例如 xarray 依赖于 numpy.fft，并将时间轴转换为频率轴，或者 Dask/xarray 内部持有 NumPy 数组以外的数组。这将在代码内部以以下方式处理：

with ua.set_backend(cupy), ua.set_backend(dask.array):
    # Code that has distributed GPU arrays here

向后兼容性

此 NEP 不提议任何向后不兼容的更改。

详细描述

提案

此 NEP 在接受时唯一提议的更改是，将 unumpy 作为覆盖 NumPy 的官方推荐方式，同时通过 uarray 默认使某些子模块可覆盖。unumpy 将仍然是一个独立的存储库/包（我们提议将其迁入以避免硬依赖，并仅在安装了单独的 unumpy 包时使用它，而不是当前依赖它）。具体来说，numpy.overridable 将成为 unumpy 的别名（如果可用），如果不可用则回退到迁入版本。uarray 和 unumpy 将主要通过通常的 GitHub 工作流程，由鸭子数组（duck-array）作者输入，其次是自定义 dtype 作者输入来开发。这样做的原因有以下几点：

• 在出现错误或问题时，迭代速度更快。

• 在需要新功能的情况下，设计更改更快。

• unumpy 也能与旧版本的 NumPy 协同工作。

• 用户和库作者选择加入覆盖过程，而不是在最不可预料的时候发生中断。简单来说，unumpy 中的错误意味着 numpy 保持不受影响。

• 对于 numpy.fft、numpy.linalg 和 numpy.random，主命名空间中的函数将镜像 numpy.overridable 命名空间中的函数。原因是这些子模块中可能存在需要后端，即使对于 numpy.ndarray 输入也是如此。

unumpy 相对于其他解决方案的优势

与为遇到的每个问题定义新协议的方法相比，unumpy 提供了许多优势：每当有需要覆盖的内容时，unumpy 都将能够提供一个统一的 API，只需进行微小的更改。例如：

• ufunc 对象可以通过它们的 __call__、reduce 等方法进行覆盖。

• 其他函数也可以以类似的方式进行覆盖。

• np.asduckarray 将消失，取而代之的是 np.overridable.asarray，并设置了后端。

• 对于像 np.zeros、np.empty 等数组创建函数也是如此。

这同样适用于未来：使某项内容可覆盖将只需要对 unumpy 进行微小的更改。

unumpy 的另一个承诺是提供默认实现。可以为任何多方法（multimethod）提供默认实现，这些实现可以基于其他多方法。这使得通过仅定义一小部分函数来轻松覆盖 NumPy API 的大部分内容。这是为了通过提供许多可以轻松地基于其他函数表达的函数的默认实现，以及一个大多数鸭子数组（duck-arrays）都需要实现的实用函数存储库，来简化新鸭子数组（duck-arrays）的创建。这将使我们避免设计完整的协议，例如，堆叠和连接协议将被简单地通过实现 stack 和/或 concatenate 来替换，然后为该类中的其他所有内容提供默认实现。这同样适用于转置，以及许多尚未提出协议的函数，例如 isin 相对于 in1d，setdiff1d 相对于 unique，等等。

它还允许以 __array_function__ 无法实现的方式覆盖函数，例如使用 opt_einsum 包中的版本覆盖 np.einsum，或者 Intel MKL 覆盖 FFT、BLAS 或 ufunc 对象。它们将定义一个具有适当多方法（multimethods）的后端，用户将通过 with 语句选择它们，或将它们注册为后端。

最后一个好处是可以通过 ua.set_backend 中的 coerce 关键字清晰地将数据转换为给定的后端，以及一个用于转换不仅数组，还有 dtype 对象和具有其他库相似的 ufunc 对象的协议。这是由于存在实际的第三方 dtype 包，以及它们希望融入 NumPy 生态系统的愿望（参见 [6]）。这与 NEP-19 中提出的 C 级 dtype 重新设计是不同的问题，它旨在允许第三方 dtype 实现与 NumPy 配合使用，就像第三方数组实现一样。这些可以提供例如单位（units）、锯齿状数组（jagged arrays）或其他超出 NumPy 范围的功能。

在同一个文件中混合 NumPy 和 unumpy

通常，人们只会想导入 unumpy 或 numpy 中的一个，出于熟悉的原因，你会将其导入为 np。然而，在某些情况下，你可能希望混合 NumPy 和覆盖，并且有几种方法可以做到这一点，具体取决于用户的风格：

from numpy import overridable as unp
import numpy as np

或

import numpy as np

# Use unumpy via np.overridable

鸭子数组（Duck-array）类型转换

从 numpy.array 或 numpy.asarray 返回非 NumPy 数组的对象存在固有的问题，特别是在 C/C++ 或 Cython 代码可能获得一个内存布局与预期不同的对象的情况下。但是，我们认为这个问题不仅适用于这两个函数，还适用于所有返回 NumPy 数组的函数。出于这个原因，覆盖是用户选择加入的，通过使用 numpy.overridable 子模块而不是 numpy。NumPy 将继续不受 numpy.overridable 中任何内容的影响而正常工作。

如果用户希望获得一个 NumPy 数组，有两种方法可以做到：

1. 使用 numpy.asarray（不可覆盖的版本）。

2. 使用 numpy.overridable.asarray 并设置了 NumPy 后端和启用了类型转换。

numpy.overridable 命名空间之外的别名

numpy.random、numpy.linalg 和 numpy.fft 中的所有功能都将别名化到 numpy.overridable 中的相应可覆盖版本。这样做的原因是存在随机数生成器（RNGs）（mkl-random）、线性代数例程（eigen, blis）和 FFT 例程（mkl-fft, pyFFTW）的替代实现，它们需要操作 numpy.ndarray 输入，但仍然需要能够切换行为。

这与 monkeypatching 在几个方面有所不同：

• 函数的面向调用者的签名始终相同，因此至少存在 API 合约的松散意义。Monkeypatching 不提供此能力。

• 能够本地切换后端。

• 有人建议 1.17 未能进入 Anaconda 默认通道的原因是 monkeypatching 和 __array_function__ 之间的不兼容性，因为 monkeypatching 会完全绕过协议。

• 诸如 from numpy import x; x 和 np.x 这样的语句，根据导入是在 monkeypatching 发生之前还是之后，结果会不同。

所有这些在 __array_function__ 或 __array_ufunc__ 中都是不可能的。

NumPy 路线图 已经正式认识到（至少部分）需要一个后端系统。

对于 numpy.random，仍然需要使 C-API 符合 NEP-19 中提出的。这对于 mkl-random 是不可能的，因为它需要重写才能适应该框架。流兼容性的保证将与以前相同，但如果存在影响 numpy.random 设置的后端，我们不对流兼容性做出任何保证，其流兼容性由后端作者自己提供。

提供隐式分发的途径

有人建议，需要能够分发不接受可分发参数的方法，同时从另一个可分发参数猜测后端。

作为一个具体的例子，考虑以下情况：

with unumpy.determine_backend(array_like, np.ndarray):
    unumpy.arange(len(array_like))

虽然这在 uarray 中尚未存在，但添加它很容易。这种代码的需要是因为人们可能希望有一个替代方案来处理提议的 *_like 函数，或者 like= 关键字参数。需要这些是因为 NumPy API 中存在不接受可分发参数的函数，但仍然需要根据另一个可分发参数来选择后端。

选择加入模块的必要性

选择加入模块的必要性是由于以下几个原因：

• API 的某些部分（如 numpy.asarray）由于与 C/Cython 扩展的兼容性问题而无法被覆盖，但是，人们可能希望通过设置了后端的 asarray 来转换为鸭子数组。

• 关于隐式选项和 monkeypatching 可能存在问题，如上所述。

NEP 18 指出这可能需要维护两个独立的 API。但是，这种负担可以通过例如通过 fixture 对 numpy.overridable 进行参数化测试来减轻。这还具有彻底测试它的副作用，不像 __array_function__。我们还认为，它提供了将 NumPy API 合约与实现正确分离的机会。

对最终用户的好处以及混合后端

在 uarray 中混合后端很容易，只需要执行：

# Explicitly say which backends you want to mix
ua.register_backend(backend1)
ua.register_backend(backend2)
ua.register_backend(backend3)

# Freely use code that mixes backends here.

对最终用户的好处不仅限于编写新代码。旧代码（通常是脚本形式）可以通过简单的导入开关和添加首选后端的一行来轻松地移植到不同的后端。这样，用户可能会发现更容易将现有代码移植到 GPU 或分布式计算。

相关工作

其他覆盖机制

• NEP-18，__array_function__ 协议。[2]

• NEP-13，__array_ufunc__ 协议。[3]

• NEP-30，__duck_array__ 协议。[9]

现有的类似 NumPy 的数组实现

• Dask: https://dask.org.cn/

• CuPy: https://cupy.chainer.org/

• PyData/Sparse: https://sparse.pydata.org.cn/

• Xnd: https://xnd.readthedocs.io/

• Astropy 的 Quantity: https://docs.astropy.org/en/stable/units/

替代数组的现有和潜在消费者

• Dask: https://dask.org.cn/

• scikit-learn: https://scikit-learn.cn/

• xarray: https://xarray.pydata.org/

• TensorLy: http://tensorly.org/

现有的替代 dtype 实现

• ndtypes: https://ndtypes.readthedocs.io/en/latest/

• Datashape: https://datashape.readthedocs.io

• Plum: https://plum-py.readthedocs.io/

NumPy API 部分的替代实现

• mkl_random: IntelPython/mkl_random

• mkl_fft: IntelPython/mkl_fft

• bottleneck: pydata/bottleneck

• opt_einsum: dgasmith/opt_einsum

实现

此 NEP 的实现将需要以下步骤：

• 在 unumpy 存储库中实现对应于 NumPy API 的 uarray 多方法（multimethods），包括用于覆盖 dtype、ufunc 和 array 对象的类，这些类通常非常容易创建。

• 将后端从 unumpy 移动到各自的数组库。

可以通过对 {numpy, unumpy} 进行参数化测试来简化维护。如果方法中添加了新参数，则需要在 unumpy 中更新相应的参数提取器和替换器。

许多参数提取器可以从 __array_function__ 协议的现有实现中重用，并且替换器通常可以在许多方法之间重用。

对于默认可覆盖的命名空间部分，主方法需要重命名并隐藏在 uarray 多方法后面。

默认实现通常在文档中使用“相当于”等词来描述，因此很容易获得。

uarray 入门

注意：本节不会试图深入介绍 uarray 的细节，这正是 uarray 文档的目的。 [1] 但是，NumPy 社区将通过 issue tracker 参与 uarray 的设计。

unumpy 是一个接口，它定义了一组与 NumPy API 兼容的可覆盖函数（多方法）。为此，它使用了 uarray 库。uarray 是一个通用工具，用于创建多方法（multimethods），这些多方法可以分发到多个不同的后端实现之一。从这个意义上说，它类似于 __array_function__ 协议，但关键区别在于后端是由最终用户显式安装的，而不是与数组类型耦合。

将后端与数组类型解耦为最终用户和后端作者提供了更大的灵活性。例如，可以：

• 覆盖不接受数组作为参数的函数

• 从数组类型创建后端

• 为同一数组类型安装多个后端

这种解耦也意味着 uarray 不仅限于分发数组类类型。后端可以自由检查整个函数参数集，以确定它是否可以实现该函数，例如 dtype 参数分发。

定义后端

uarray 由两个主要协议组成：__ua_convert__ 和 __ua_function__，按此顺序调用，以及 __ua_domain__。__ua_convert__ 用于转换和类型转换。它的签名是 (dispatchables, coerce)，其中 dispatchables 是 ua.Dispatchable 对象的迭代器，coerce 是一个布尔值，指示是否强制转换。ua.Dispatchable 是一个简单的类，由三个简单值组成：type、value 和 coercible。__ua_convert__ 返回转换后的值的迭代器，如果失败则返回 NotImplemented。

__ua_function__ 的签名是 (func, args, kwargs)，它定义了函数的实际实现。它接收函数及其参数。返回 NotImplemented 将导致移动到函数的默认实现（如果存在），如果不存在，则移动到下一个后端。

假设调用了一个 uarray 多方法，以下是会发生的情况：

1. 我们规范化参数，以便任何没有默认值的参数都放在 *args 中，有默认值的参数放在 **kwargs 中。

2. 我们检查后端列表。

   1. 如果列表为空，我们尝试默认实现。

3. 我们检查后端的 __ua_convert__ 方法是否存在。如果存在：

   1. 我们将分发器的输出（这是一个 ua.Dispatchable 对象的迭代器）传递给它。

   2. 我们将此输出与参数一起馈送到参数替换器。NotImplemented 表示我们使用下一个后端移动到 3。

   3. 我们将替换后的参数存储为新参数。

4. 我们将参数馈送到 __ua_function__，并返回输出，如果输出不是 NotImplemented，则退出。

5. 如果默认实现存在，我们则尝试使用当前后端执行它。

6. 失败时，我们使用下一个后端移动到 3。如果不再有后端，我们移动到 7。

7. 我们引发一个 ua.BackendNotImplementedError。

定义可覆盖的多方法

要定义一个可覆盖的函数（多方法），需要以下几项：

1. 一个返回 ua.Dispatchable 对象迭代器的分发器。

2. 一个替换可分发值（dispatchable values）为所提供值的反向分发器。

3. 一个域（domain）。

4. 可选地，一个默认实现，可以用其他多方法来提供。

例如，考虑以下情况：

import uarray as ua

def full_argreplacer(args, kwargs, dispatchables):
    def full(shape, fill_value, dtype=None, order='C'):
        return (shape, fill_value), dict(
            dtype=dispatchables[0],
            order=order
        )

    return full(*args, **kwargs)

@ua.create_multimethod(full_argreplacer, domain="numpy")
def full(shape, fill_value, dtype=None, order='C'):
    return (ua.Dispatchable(dtype, np.dtype),)

在 unumpy 存储库 [8] 中可以找到大量示例。这种简单的覆盖可调用对象（callables）的行为使我们能够覆盖：

• 方法

• 属性，通过 fget 和 fset。

• 整个对象，通过 __get__。

NumPy 的示例

实现 NumPy 类似 API 的库将以以下方式使用它（作为示例）：

import numpy.overridable as unp
_ua_implementations = {}

__ua_domain__ = "numpy"

def __ua_function__(func, args, kwargs):
    fn = _ua_implementations.get(func, None)
    return fn(*args, **kwargs) if fn is not None else NotImplemented

def implements(ua_func):
    def inner(func):
        _ua_implementations[ua_func] = func
        return func

    return inner

@implements(unp.asarray)
def asarray(a, dtype=None, order=None):
    # Code here
    # Either this method or __ua_convert__ must
    # return NotImplemented for unsupported types,
    # Or they shouldn't be marked as dispatchable.

# Provides a default implementation for ones and zeros.
@implements(unp.full)
def full(shape, fill_value, dtype=None, order='C'):
    # Code here

替代方案

解决此问题的当前替代方案是 NEP-18 [2]、NEP-13 [4] 和 NEP-30 [9] 的组合，再加上添加更多协议（尚未指定）。即使那样，NumPy API 的某些部分仍将是不可覆盖的，因此它是一个部分替代方案。

迁入 unumpy 的主要替代方案是简单地将其完全移入 NumPy，而不是作为单独的包分发。这也将实现提议的目标，但我们更愿意暂时将其保留为一个单独的包，原因已在上面说明。

第三个替代方案是将 unumpy 移到 NumPy 组织下，并作为 NumPy 项目进行开发。这也将实现所述目标，并且也是此 NEP 可以考虑的一种可能性。但是，执行额外的 pip install 或 conda install 可能会劝退一些用户采用此方法。

要求选择加入的替代方案主要是不重写 np.asarray 和 np.array，并使 NumPy API 的其余部分可重写，而是提供 np.duckarray 和 np.asduckarray 作为使用相应重写的、对 duck-array 友好的替代方案。然而，这样做会增加 NumPy 调用的一些微小开销。

讨论

• uarray 博客文章： https://labs.quansight.org/blog/2019/07/uarray-update-api-changes-overhead-and-comparison-to-__array_function__/

• NEP 18 — A dispatch mechanism for NumPy’s high level array functions 的讨论部分

• NEP 22 — NumPy 数组的鸭子类型 - 高层概览

• Dask issue #4462： dask/dask#4462

• PR #13046： numpy/numpy#13046

• Dask issue #4883： dask/dask#4883

• Issue #13831： numpy/numpy#13831

• 讨论 PR 1： hameerabbasi/numpy#3

• 讨论 PR 2： hameerabbasi/numpy#4

• 讨论 PR 3： numpy/numpy#14389

参考文献和脚注

[1]

uarray，Python 的通用分派机制： https://uarray.readthedocs.io

[2] (1,2,3,4)

NEP 18 — NumPy 高层数组函数的派发机制

[3] (1,2)

NEP 22 — NumPy 数组的鸭子类型 - 高层概览

[4] (1,2)

NEP 13 — 覆盖 UFunc 的机制

[5]

回复“为 NumPy 方法的非分派实现添加内容”： https://mail.python.org/archives/list/numpy-discussion@python.org/thread/5GUDMALWDIRHITG5YUOCV343J66QSX3U/#5GUDMALWDIRHITG5YUOCV343J66QSX3U

[6]

自定义 Dtype/Units 讨论： https://mail.python.org/archives/list/numpy-discussion@python.org/thread/RZYCVT6C3F7UDV6NA6FEV4MC5FKS6RDA/#RZYCVT6C3F7UDV6NA6FEV4MC5FKS6RDA

[7]

史诗级的 dtype 清理计划： numpy/numpy#2899

[8] (1,2)

unumpy：NumPy，但实现无关： https://unumpy.readthedocs.io

[9] (1,2,3,4)

NEP 30 — NumPy 数组的鸭子类型 - 实现

[10]

http://scipy.github.io/doc/stable/fft.html#backend-control

版权

本文档已置于公共领域。