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/[email protected]/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 中进行覆盖而无需设计新的协议。

uarray 背后的动机是多方面的：首先，已经进行了几次尝试以允许调度 NumPy API 的部分，包括（最突出的是）NEP-13 [4] 中的 __array_ufunc__ 协议和 NEP-18 [2] 中的 __array_function__ 协议，但这表明需要开发更多协议，包括强制转换协议（参见 [5]，[9]）。在参考文献中已经广泛讨论了需要这些覆盖的原因，本 NEP 不会尝试详细说明为什么需要这些；但简而言之：库作者必须能够将任意对象强制转换为他们自己类型的数组，例如 CuPy 需要强制转换为 CuPy 数组，而不是 NumPy 数组。简单来说，需要像 np.asarray(...) 这样的东西或其他方法才能“正常工作”并返回鸭子数组。

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

用法和影响

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

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

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

# 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的主要开发人员将是duck-array的作者，其次是自定义dtype的作者，通过通常的GitHub工作流程进行。这样做有几个原因：

• 在出现bug或问题时，可以更快地迭代。

• 在需要功能的情况下，可以更快地进行设计更改。

• unumpy也能与旧版本的NumPy一起工作。

• 用户和库作者可以选择是否进行覆盖，而不是在最意想不到的时候发生中断。简单来说，unumpy中的bug不会影响numpy。

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

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

unumpy相对于为遇到的每个问题定义一个新协议的方法，具有许多优势：每当需要覆盖时，unumpy都能提供具有微小更改的统一API。例如：

• 可以通过__call__、reduce和其他方法覆盖ufunc对象。

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

• np.asduckarray将被移除，并变为具有后端设置的np.overridable.asarray。

• 数组创建函数（如np.zeros、np.empty等）也是如此。

这同样适用于未来：使某些内容可覆盖只需要对unumpy进行少量更改。

unumpy的另一个优势是默认实现。可以为任何多方法提供默认实现，以其他方法表示。这允许通过仅定义一小部分NumPy API来覆盖大部分NumPy API。这简化了新duck-array的创建，方法是为许多可以用其他方法轻松表示的函数提供默认实现，以及提供大多数duck-array所需的实用函数的存储库。这将使我们能够避免设计整个协议，例如，堆叠和连接的协议将被简单地实现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

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内的相应可覆盖版本关联。原因是存在需要对numpy.ndarray输入进行操作的RNG（mkl-random）、线性代数例程（eigen、blis）和FFT例程（mkl-fft、pyFFTW）的替代实现，但仍然需要能够切换行为。

这与猴子补丁在以下几个方面有所不同：

• 函数的调用者可见签名始终相同，因此至少存在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将数据强制转换为duck-array。

• 隐式选项和猴子补丁可能存在一些问题，如上所述。

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/

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

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

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

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

现有和潜在的替代数组使用者

• Dask: https://dask.org/

• 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对象的iterable，coerce是一个布尔值，指示是否强制转换。ua.Dispatchable是一个简单的类，包含三个简单的值：type、value和coercible。__ua_convert__返回转换后的值的iterable，或者在失败的情况下返回NotImplemented。

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

假设调用了uarray多方法，将会发生以下情况：

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

2. 我们检查后端列表。

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

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

   1. 我们将分派器的输出传递给它，这是一个ua.Dispatchable对象的iterable。

   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 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/[email protected]/thread/5GUDMALWDIRHITG5YUOCV343J66QSX3U/#5GUDMALWDIRHITG5YUOCV343J66QSX3U

[6]

自定义 Dtype/Units 讨论：https://mail.python.org/archives/list/[email protected]/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/devdocs/fft.html#backend-control

版权

本文档已进入公有领域。