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 提议以选择加入的方式，允许覆盖 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] 成为 NumPy API 中尚未被 __array_function__ 或 __array_ufunc__ 覆盖的部分的推荐覆盖机制，并且 uarray 被引入 NumPy 内的一个新命名空间，以便用户和下游依赖项能够访问这些覆盖。这种引入机制类似于 SciPy 为使 scipy.fft 可覆盖而采取的做法（参见 [10]）。

uarray 背后的动机是多方面的：首先，已经有几次尝试允许调度 NumPy API 的部分功能，包括（最突出的是）NEP-13 [4] 中的 __array_ufunc__ 协议，以及 NEP-18 [2] 中的 __array_function__ 协议，但这表明需要开发更多协议，包括一个用于强制转换的协议（参见 [5], [9]）。需要这些覆盖的原因已在参考文献中 extensively 讨论，本 NEP 将不试图深入探讨为什么需要这些；但简而言之：对于库作者来说，有必要能够将任意对象强制转换为他们自己类型的数组，例如，CuPy 需要强制转换为 CuPy 数组，而不是 NumPy 数组。简单来说，需要像 np.asarray(...) 这样的东西，或者一个能够“直接工作”并返回鸭子数组的替代方案。

用法和影响

本 NEP 允许全局和语境本地覆盖，以及类似 __array_function__ 的自动覆盖。

除了动机部分中提到的第一个用例之外，本 NEP 将支持以下一些用例

第一个是允许替代数据类型返回其各自的数组。

# 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)

这将允许以官方方式使覆盖与 NumPy 配合工作，而无需猴子补丁或分发修改过的 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 将主要根据鸭子数组作者的输入进行开发，其次根据自定义 dtype 作者的输入进行开发，通过常见的 GitHub 工作流程。这有几个原因

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

• 在需要功能时，更快的设计变更。

• 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 的另一个承诺是提供默认实现。任何多方法都可以提供默认实现，这些实现可以基于其他多方法。这允许通过仅定义 NumPy API 的一小部分来覆盖其大部分功能。这是为了简化新鸭子数组的创建，通过提供许多可以轻松地用其他函数表达的函数的默认实现，以及一个包含实用函数的仓库，这些函数有助于实现大多数鸭子数组所需的鸭子数组。这将使我们能够避免设计整个协议，例如，堆叠和连接的协议将被简单地通过实现 stack 和/或 concatenate，然后为该类中的所有其他内容提供默认实现来取代。同样适用于转置，以及许多尚未提出协议的其他函数，例如通过 in1d 实现 isin，通过 unique 实现 setdiff1d，等等。

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

最后一个好处是提供了一种明确的方式来强制转换到给定的后端（通过 ua.set_backend 中的 coerce 关键字），以及一个协议，用于强制转换不仅是数组，还包括来自其他库的相似的 dtype 对象和 ufunc 对象。这是由于实际存在第三方 dtype 包，以及它们融入 NumPy 生态系统的愿望（参见 [6]）。与 [7] 中提议的 C 级 dtype 重新设计相比，这是一个单独的问题，它关注的是允许第三方 dtype 实现与 NumPy 协同工作，就像第三方数组实现一样。这些可以提供例如单位、锯齿数组或其他超出 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

鸭子数组强制转换

从 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 内部各自可覆盖的版本。这样做的原因是因为存在随机数生成器（mkl-random）、线性代数例程（eigen、blis）和 FFT 例程（mkl-fft、pyFFTW）的替代实现，它们需要对 numpy.ndarray 输入进行操作，但仍需要切换行为的能力。

这在几个方面与猴子补丁不同

• 函数的面向调用者签名始终相同，因此至少存在一种松散的 API 契约感。猴子补丁不提供此功能。

• 具有局部切换后端的能力。

• 有人提出，1.17 版尚未进入 Anaconda 默认通道的原因是猴子补丁和 __array_function__ 之间存在不兼容性，因为猴子补丁会完全绕过该协议。

• 形式为 from numpy import x; x 和 np.x 的语句会产生不同的结果，具体取决于导入是在猴子补丁发生之前还是之后进行。

所有这些都无法通过 __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 强制转换为鸭子数组。

• 隐式选项和猴子补丁可能存在问题，例如上面提到的那些。

NEP 18 指出这可能需要维护两个独立的 API。然而，例如，可以通过夹具分别对 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/

• 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 多方法，包括用于覆盖 dtype、ufunc 和 array 对象的类，这些通常非常容易创建。

• 将后端从 unumpy 移入各自的数组库中。

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

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

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

默认实现在文档中通常使用“等价于”字样，因此，很容易获取。

uarray 入门

注意： 本节不会试图过于详细地介绍 uarray，那是 uarray 文档的目的。 [1] 然而，NumPy 社区将通过问题跟踪器对 uarray 的设计提供意见。

unumpy 是一个接口，它定义了一组与 NumPy API 兼容的可覆盖函数（多方法）。为此，它使用了 uarray 库。uarray 是一个通用工具，用于创建可调度到多个不同可能后端实现之一的多方法。从这个意义上讲，它类似于 __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. 一个反向调度器，它用提供的值替换可调度值。

3. 一个领域。

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]。这种简单的可调用对象覆盖行为允许我们覆盖

• 方法

• 属性，通过 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 作为使用相应覆盖的鸭子数组友好替代方案。然而，这有一个缺点，即会给 NumPy 调用增加少量开销。

讨论

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

• NEP 18 — NumPy 高级数组函数的调度机制 的讨论部分

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

• Dask 问题 #4462：dask/dask#4462

• PR #13046：numpy/numpy#13046

• Dask 问题 #4883：dask/dask#4883

• 问题 #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

版权

本文档已置于公共领域。