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.zerosnumpy.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.arraycupy 本身。

这些类型的覆盖对于最终用户和库作者都很有用。最终用户可能已经编写或希望编写代码,然后他们稍后可以加速或迁移到不同的实现,例如 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 [4] 中的 __array_ufunc__ 协议和 NEP-18 [2] 中的 __array_function__ 协议,但这表明需要开发更多协议,包括强制转换协议(参见 [5][9])。这些覆盖需要的原因在参考文献中已广泛讨论,本 NEP 不会尝试详细说明为什么需要这些覆盖;但简而言之:库作者有必要能够将任意对象强制转换为他们自己的类型的数组,例如 CuPy 需要强制转换为 CuPy 数组,而不是 NumPy 数组。简单来说,需要像 np.asarray(...) 或其他替代方案一样“正常工作”并返回鸭子数组。

用法和影响#

本 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 的别名(如果可用),如果不可用则回退到引入的版本。 uarrayunumpy 将主要在鸭子数组作者的输入下开发,其次是自定义 dtype 作者,通过通常的 GitHub 工作流程。这样做的原因有几个

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

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

  • unumpy 也适用于旧版本的 NumPy。

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

  • 对于 numpy.fftnumpy.linalgnumpy.random,主命名空间中的函数将反映 numpy.overridable 命名空间中的函数。这样做的原因是,即使对于 numpy.ndarray 输入,这些子模块中也可能存在需要后端的函数。

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

unumpy 相对于为遇到的每个问题定义新协议的方法提供了许多优势:每当有需要覆盖的内容时,unumpy 都能够提供具有非常少量更改的统一 API。例如

  • ufunc 对象可以通过它们的 __call__reduce 和其他方法进行覆盖。

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

  • np.asduckarray 不再存在,并成为具有已设置后端的 np.overridable.asarray

  • 数组创建函数(例如 np.zerosnp.empty 等)也是如此。

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

unumpy 另一个承诺是默认实现。可以针对任何多方法以其他方法的形式提供默认实现。这允许通过仅定义其一小部分来覆盖 NumPy API 的大部分内容。这是为了简化新鸭子数组的创建,方法是为许多可以用其他方法轻松表达的函数提供默认实现,以及一个帮助实现大多数鸭子数组所需的实用函数的存储库。这将使我们能够避免设计整个协议,例如,堆叠和连接的协议将被简单地实现 stack 和/或 concatenate 然后为该类中的所有其他内容提供默认实现所取代。转置也是如此,以及许多其他尚未提出协议的函数,例如 isin 就以 in1d 的形式存在,setdiff1d 就以 unique 的形式存在,等等。

它还允许以__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#

通常,人们只希望导入unumpynumpy中的一个,为了熟悉起见,你会将其导入为np。但是,在某些情况下,人们可能希望混合使用NumPy和覆盖,并且有几种方法可以做到这一点,具体取决于用户的风格

from numpy import overridable as unp
import numpy as np

import numpy as np

# Use unumpy via np.overridable

Duck数组强制转换#

numpy.arraynumpy.asarray返回不是NumPy数组的对象存在固有问题,尤其是在C/C++或Cython代码的上下文中,这些代码可能会获得一个与预期内存布局不同的对象。但是,我们认为这个问题可能不仅适用于这两个函数,还适用于所有返回NumPy数组的函数。因此,覆盖是用户可选的,通过使用子模块numpy.overridable而不是numpy。NumPy将继续不受numpy.overridable中任何内容的影响。

如果用户希望获得NumPy数组,则有两种方法可以实现

  1. 使用numpy.asarray(不可覆盖版本)。

  2. 在设置NumPy后端并启用强制转换的情况下,使用numpy.overridable.asarray

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

numpy.randomnumpy.linalgnumpy.fft中的所有功能都将在numpy.overridable内部映射到它们各自的可覆盖版本。这样做的原因是存在RNG的替代实现(mkl-random)、线性代数例程(eigenblis)和FFT例程(mkl-fftpyFFTW),它们需要对numpy.ndarray输入进行操作,但仍然需要能够切换行为。

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

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

  • 能够在本地切换后端。

  • 有人建议,1.17没有在Anaconda默认通道中发布的原因是猴子补丁与__array_function__之间的不兼容性,因为猴子补丁将完全绕过该协议。

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

所有这些都无法使用__array_function____array_ufunc__实现。

NumPy路线图中,至少在某种程度上已经正式认识到需要一个后端系统来实现这一点。

对于numpy.random,仍然需要使C-API适合NEP-19中提出的C-API。对于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数组。

  • 围绕隐式选项和猴子补丁可能会出现一些问题,例如上面提到的那些问题。

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的实施将需要以下步骤

  • unumpy存储库中实现与NumPy API兼容的uarray多方法,包括用于覆盖dtypeufuncarray对象的类,这些类通常非常易于创建。

  • 将后端从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),其中dispatchablesua.Dispatchable对象的iterable,coerce是指示是否强制转换的布尔值。ua.Dispatchable是一个简单的类,包含三个简单值:typevaluecoercible__ua_convert__返回转换后的值的iterable,或者在失败的情况下返回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]。这种简单的重写可调用对象的行为允许我们重写

  • 方法

  • 属性,通过 fgetfset

  • 整个对象,通过 __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 installconda install 操作可能会阻止某些用户采用此方法。

不需要选择加入的主要替代方案是*不*重写 np.asarraynp.array,而是使 NumPy API 的其余部分表面可重写,而是提供 np.duckarraynp.asduckarray 作为鸭子数组友好的替代方案,这些方案使用相应的重写。但是,这样做会增加 NumPy 调用的少量开销。

讨论#

参考文献和脚注#