NEP 37 — 面向类 NumPy 模块的调度协议#

作者:

Stephan Hoyer <shoyer@google.com>

作者:

Hameer Abbasi

作者:

Sebastian Berg

状态:

已取代

取代者:

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

类型:

标准跟踪

创建时间:

2019-12-29

决议:

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

摘要#

NEP-18 的 __array_function__ 喜忧参半。一些项目(例如 dask、CuPy、xarray、sparse、Pint、MXNet)积极采纳了它,而另一些项目(例如 JAX)则更不情愿。在此,我们提出了一种新的协议 __array_module__,我们预计它最终可以取代 __array_function__ 的大多数用例。该协议需要用户和库作者的明确采用,这确保了向后兼容性,并且它也比 __array_function__ 简单得多,我们预计这两点将使其更容易被采纳。

为什么 __array_function__ 仍不足够#

NEP-18 未能实现其目标主要有两个方面

  1. 向后兼容性问题__array_function__ 对使用它的库具有重大影响

    • JAX 不愿实现 __array_function__ 的部分原因是它担心破坏现有代码:用户期望像 np.concatenate 这样的 NumPy 函数返回 NumPy 数组。这是 __array_function__ 设计的一个基本限制,我们选择它来允许覆盖现有的 numpy 命名空间。像 Dask 和 CuPy 这样的库已经考虑并接受了 __array_function__ 的向后不兼容性影响;如果这种影响不存在,对它们来说仍然会更好。

      请注意,像 PyTorchscipy.sparse 这样的项目也尚未采用 __array_function__,因为它们没有 NumPy 兼容的 API 或语义。对于 PyTorch,这在未来可能会添加。scipy.sparse 的情况与 numpy.matrix 类似:其语义与 numpy.ndarray 不兼容,因此添加 __array_function__(除非返回 NotImplemented)不是一个好主意。

    • __array_function__ 目前要求采用“全有或全无”的方法来实现 NumPy 的 API。没有良好的 增量采用 途径,这对于已有的项目来说尤其成问题,因为采用 __array_function__ 将导致破坏性更改。

  2. 可覆盖内容的限制。 __array_function__ 存在一些重要空白,最显著的是数组创建和强制转换函数

    • 数组创建 例程(例如 np.arangenp.random 中的例程)需要其他机制来指示要创建的数组类型。NEP 35 提议为没有现有数组参数的函数添加可选的 like= 参数。然而,我们仍然缺乏任何覆盖对象方法(例如 np.random.RandomState 所需的方法)的机制。

    • 数组转换 无法重用现有的强制转换函数,例如 np.asarray,因为 np.asarray 有时表示“转换为精确的 np.ndarray”,而其他时候表示“转换为_类似_ NumPy 数组的东西”。这导致了 NEP 30 提案,即一个独立的 np.duckarray 函数,但这仍然无法解决如何将一个鸭子数组转换为与另一个鸭子数组匹配的类型。

其他提出的可维护性问题包括

  • 在支持覆盖的模块中,不再可能使用 NumPy 函数的别名。例如,CuPy 和 JAX 都将 result_type = np.result_type,现在必须将 np.result_type 的使用包装到它们自己的 result_type 函数中。

  • 为未实现的 NumPy 函数实现 回退机制 (通过使用 NumPy 的实现)很难做到正确(但请参阅 dask 的版本),因为 __array_function__ 不提供一致的接口。转换所有数组类型的参数需要递归处理 *args, **kwargs 形式的通用参数。

get_array_module__array_module__ 协议#

我们提出了一种新的面向用户的机制,用于调度到鸭子数组实现:numpy.get_array_moduleget_array_module 执行与 __array_function__ 相同的类型解析,并返回一个其 API 承诺与 numpy 标准接口匹配的模块,该模块可以实现对所有提供的数组类型的操作。

该协议本身比 __array_function__ 更简单且更强大,因为它不需要担心实际实现函数。我们相信它解决了 __array_function__ 的大部分可维护性和功能限制。

新协议是可选加入、显式且具有局部控制的;有关这些设计特征重要性的讨论,请参阅 附录:API 覆盖的设计选择

数组模块约定#

get_array_module/__array_module__ 返回的模块应尽最大努力在新数组类型上实现 NumPy 的核心功能。未实现的功能应直接省略(例如,访问未实现的功能应引发 AttributeError)。未来,我们预计将制定一个协议,用于请求 numpy 的受限子集;更多详情请参阅 请求 NumPy API 的受限子集

如何使用 get_array_module#

希望支持通用鸭子数组的代码应显式调用 get_array_module 来确定要从中调用函数的适当数组模块,而不是直接使用 numpy 命名空间。例如

# calls the appropriate version of np.something for x and y
module = np.get_array_module(x, y)
module.something(x, y)

数组创建和数组转换都受支持,因为调度由 get_array_module 处理,而不是通过函数参数的类型。例如,要使用随机数生成函数或方法,我们只需取出相应的子模块

def duckarray_add_random(array):
    module = np.get_array_module(array)
    noise = module.random.randn(*array.shape)
    return array + noise

我们还可以编写 NEP 30 中的鸭子数组 stack 函数,而无需新的 np.duckarray 函数

def duckarray_stack(arrays):
    module = np.get_array_module(*arrays)
    arrays = [module.asarray(arr) for arr in arrays]
    shapes = {arr.shape for arr in arrays}
    if len(shapes) != 1:
        raise ValueError('all input arrays must have the same shape')
    expanded_arrays = [arr[module.newaxis, ...] for arr in arrays]
    return module.concatenate(expanded_arrays, axis=0)

默认情况下,如果没有任何参数是数组,get_array_module 将返回 numpy 模块。可以通过提供仅关键字参数 module 来显式控制此回退。还可以通过设置 module=None 来指示应引发异常而不是返回默认数组模块。

如何实现 __array_module__#

实现鸭子数组类型并希望支持 get_array_module 的库需要实现相应的协议 __array_module__。这个新协议基于 Python 的算术调度协议,本质上是 __array_function__ 的简化版本。

__array_module__ 只传递一个参数,即传入 get_array_module 的唯一数组类型的 Python 集合,换句话说,所有具有 __array_module__ 属性的参数。

这个特殊方法应该返回一个 API 与 numpy 匹配的命名空间,或者返回 NotImplemented,表示它不知道如何处理该操作。

class MyArray:
    def __array_module__(self, types):
        if not all(issubclass(t, MyArray) for t in types):
            return NotImplemented
        return my_array_module

__array_module__ 返回自定义对象#

my_array_module 通常(但不总是)会是一个 Python 模块。返回自定义对象(例如,通过 __getattr__ 实现函数)可能对某些高级用例有用。

例如,自定义对象可以允许鸭子数组模块的部分实现,并回退到 NumPy(尽管通常不推荐这样做,因为此类回退行为容易出错)

class MyArray:
    def __array_module__(self, types):
        if all(issubclass(t, MyArray) for t in types):
            return ArrayModule()
        else:
            return NotImplemented

class ArrayModule:
    def __getattr__(self, name):
        import base_module
        return getattr(base_module, name, getattr(numpy, name))

numpy.ndarray 继承子类#

NEP-18 中关于明确定义类型转换层次结构的所有相同指导仍然适用。numpy.ndarray 本身包含一个匹配的 __array_module__ 实现,这对子类很方便。

class ndarray:
    def __array_module__(self, types):
        if all(issubclass(t, ndarray) for t in types):
            return numpy
        else:
            return NotImplemented

NumPy 的内部机制#

get_array_module 的类型解析规则遵循 Python 和 NumPy 现有调度协议的相同模型:子类在超类之前被调用,否则从左到右。__array_module__ 保证在每个唯一类型上只被调用一次。

get_array_module 的实际实现将使用 C 语言,但应与以下 Python 代码等效

def get_array_module(*arrays, default=numpy):
    implementing_arrays, types = _implementing_arrays_and_types(arrays)
    if not implementing_arrays and default is not None:
        return default
    for array in implementing_arrays:
        module = array.__array_module__(types)
        if module is not NotImplemented:
            return module
    raise TypeError("no common array module found")

def _implementing_arrays_and_types(relevant_arrays):
    types = []
    implementing_arrays = []
    for array in relevant_arrays:
        t = type(array)
        if t not in types and hasattr(t, '__array_module__'):
            types.append(t)
            # Subclasses before superclasses, otherwise left to right
            index = len(implementing_arrays)
            for i, old_array in enumerate(implementing_arrays):
                if issubclass(t, type(old_array)):
                    index = i
                    break
            implementing_arrays.insert(index, array)
    return implementing_arrays, types

__array_ufunc____array_function__ 的关系#

这些旧协议具有不同的用例,应保留#

__array_module__ 旨在解决 __array_function__ 的限制,因此自然会考虑它是否可以完全取代 __array_function__。这将带来双重好处:(1) 简化用户关于如何覆盖 NumPy 的叙述,以及 (2) 消除每次调用 NumPy 函数时检查调度所带来的减速。

然而,从用户角度来看,__array_module____array_function__ 相当不同:前者需要显式调用 get_array_function,而不是简单地重用原始的 numpy 函数。这对于依赖鸭子数组的 来说可能没问题,但对于交互式使用来说可能过于冗长,令人沮丧。

__array_ufunc__ 的一些调度用例也可以通过 __array_module__ 解决,但并非全部。例如,仍然可以以通用方式在非 NumPy 数组(例如,使用 dask.array)上定义非 NumPy ufuncs(例如,来自 Numba 或 SciPy 的 ufuncs),这仍然很有用。

考虑到它们现有的采纳和不同的用例,我们认为目前没有理由删除或废弃 __array_function____array_ufunc__

用于实现 __array_function____array_ufunc__ 的 Mixin 类#

尽管面向用户存在差异,但 __array_module__ 和实现 NumPy API 的模块仍然包含足够的必要功能,以使用现有鸭子数组协议实现调度。

例如,以下混入(mixin)类将为这些特殊方法提供基于 get_array_module__array_module__ 的合理默认值

class ArrayUfuncFromModuleMixin:

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        arrays = inputs + kwargs.get('out', ())
        try:
            array_module = np.get_array_module(*arrays)
        except TypeError:
            return NotImplemented

        try:
            # Note this may have false positive matches, if ufunc.__name__
            # matches the name of a ufunc defined by NumPy. Unfortunately
            # there is no way to determine in which module a ufunc was
            # defined.
            new_ufunc = getattr(array_module, ufunc.__name__)
        except AttributeError:
            return NotImplemented

        try:
            callable = getattr(new_ufunc, method)
        except AttributeError:
            return NotImplemented

        return callable(*inputs, **kwargs)

class ArrayFunctionFromModuleMixin:

    def __array_function__(self, func, types, args, kwargs):
        array_module = self.__array_module__(types)
        if array_module is NotImplemented:
            return NotImplemented

        # Traverse submodules to find the appropriate function
        modules = func.__module__.split('.')
        assert modules[0] == 'numpy'
        for submodule in modules[1:]:
            module = getattr(module, submodule, None)
        new_func = getattr(module, func.__name__, None)
        if new_func is None:
            return NotImplemented

        return new_func(*args, **kwargs)

为了更方便编写鸭子数组,我们也可以将这些混入类添加到 numpy.lib.mixins 中(但上述示例可能已足够)。

考虑的替代方案#

命名#

我们喜欢 __array_module__ 这个名称,因为它与现有的 __array_function____array_ufunc__ 协议相呼应。另一个合理的选择可能是 __array_namespace__

调用此协议的 NumPy 函数应该如何命名尚不明确(本提案中为 get_array_module)。一些可能的替代方案包括:array_modulecommon_array_moduleresolve_array_moduleget_namespaceget_numpyget_numpylike_moduleget_duck_array_module

请求 NumPy API 的受限子集#

随着时间的推移,NumPy 已经积累了一个非常庞大的 API 表面,仅顶级 numpy 模块中就有超过 600 个属性。任何鸭子数组库都不太可能或不愿意实现所有这些函数和类,因为 NumPy 常用子集要小得多。

我们认为定义 NumPy API 的“最小”子集,省略很少使用或不推荐的功能,将是一项有益的实践。例如,最小的 NumPy 可能包括 stack,但不包括其他堆叠函数 column_stackdstackhstackvstack。这可以清楚地向鸭子数组的作者和用户指明哪些功能是核心的,哪些功能可以跳过。

支持请求 NumPy API 的受限子集将是 get_array_function__array_module__ 中一个自然而然的功能,例如:

# array_module is only guaranteed to contain "minimal" NumPy
array_module = np.get_array_module(*arrays, request='minimal')

为了方便与 NumPy 进行测试以及与任何有效的鸭子数组库一起使用,当 get_array_module 仅在 NumPy 数组上调用时,NumPy 本身将返回 numpy 模块的受限版本。被省略的函数将根本不存在。

不幸的是,我们尚未确定这些受限子集应该是什么,所以目前这样做没有意义。如果将来我们确定了,我们可以选择向 get_array_module 添加新的关键字参数,或者添加新的顶级函数,例如 get_minimal_array_module。我们还需要添加一个新的协议,其模式模仿 __array_module__(例如 __array_module_minimal__),或者可以向 __array_module__ 添加一个可选的第二个参数(用 try/except 捕获错误)。

用于隐式调度的新的命名空间#

我们不是在主 numpy 命名空间中通过 __array_function__ 支持覆盖,而是可以创建一个新的可选加入命名空间,例如 numpy.api,其中包含支持调度的 NumPy 函数版本。这些覆盖将需要新的可选加入协议,例如模仿 __array_function____array_function_api__

这将通过可选加入的方式解决 __array_function__ 的最大限制,并且还可以明确地覆盖像 asarray 这样的函数,因为 np.api.asarray 将始终意味着“转换一个类似数组的对象”。但它无法解决 __array_module__ 所满足的所有调度需求,并且会让我们为数组用户和实现者支持一个复杂得多的协议。

我们有可能 通过 __array_module__ 协议实现这样一个新命名空间。当然,一些用户会觉得这很方便,因为它稍微减少了样板代码。但这会给用户留下一个令人困惑的选择:他们何时应该使用 get_array_modulenp.api.something。此外,我们必须添加和维护一个全新的模块,这比仅仅添加一个函数要昂贵得多。

基于类型和数组进行调度而非仅基于类型#

我们除了仅通过唯一数组类型支持调度外,还可以通过数组对象支持调度,例如,通过在 __array_module__ 协议中传递一个 arrays 参数。这可能对带有元数据的数组(如 Dask 和 Pint 提供的数组)的调度很有用,但会在类型安全和复杂性方面带来成本。

例如,一个支持 CPU 和 GPU 数组的库可能会根据输入参数决定从 ones 等函数创建新数组的设备

class Array:
    def __array_module__(self, types, arrays):
        useful_arrays = tuple(a in arrays if isinstance(a, Array))
        if not useful_arrays:
            return NotImplemented
        prefer_gpu = any(a.prefer_gpu for a in useful_arrays)
        return ArrayModule(prefer_gpu)

class ArrayModule:
    def __init__(self, prefer_gpu):
        self.prefer_gpu = prefer_gpu

    def __getattr__(self, name):
        import base_module
        base_func = getattr(base_module, name)
        return functools.partial(base_func, prefer_gpu=self.prefer_gpu)

这可能很有用,但我们是否真的需要它尚不清楚。Pint 似乎在没有任何显式数组创建例程的情况下运行良好(倾向于通过单位进行乘法,例如 np.ones(5) * ureg.m),而且 Dask 大部分也接受现有的 __array_function__ 风格覆盖(例如,倾向于 np.ones_like 而非 np.ones)。选择将数组放置在 CPU 还是 GPU 上可以通过使数组创建惰性化来解决。

附录:API 覆盖的设计选择#

覆盖 NumPy API 的设计选择范围很广。这里我们讨论指导我们设计 __array_module__ 的三个主要设计决策轴心。

用户的可选加入与可选退出#

__array_ufunc____array_function__ 协议提供了一种在 NumPy 现有命名空间内 覆盖 NumPy 函数的机制。这意味着如果用户不希望出现任何被覆盖的行为,他们需要显式选择退出,例如通过使用 np.asarray() 转换数组。

理论上,这种方法降低了用户代码和库采用这些协议的门槛,因为使用标准 NumPy 命名空间的代码会自动兼容。但实际上,这并未奏效。例如,大多数维护良好的使用 NumPy 的库都遵循最佳实践,即使用 np.asarray() 转换所有输入,而要使用 __array_function__,它们就必须显式放宽此要求。我们的经验是,使一个库兼容新的鸭子数组类型通常至少需要少量工作来适应数据模型和可高效实现的操作之间的差异。

这些可选退出方法也大大增加了采用这些协议的库的向后兼容性,因为作为库选择加入意味着它们的用户也随之选择加入,无论用户是否预料到。对于那些未能采用 __array_function__ 的库来说,可选加入的方法似乎是必不可少的。

显式与隐式实现选择#

__array_ufunc____array_function__ 都对调度具有隐式控制:每次函数调用时,调度函数通过相应的协议确定。这很好地概括了处理许多不同类型对象的情况,Python 中算术运算符的实现就证明了这一点,但它对可读性有一个重要的缺点:代码读者不再能立即清楚地知道函数被调用时发生了什么,因为函数的实现可能被其任何参数覆盖。

速度方面的影响是

  • 当使用 鸭子数组类型 时,get_array_module 意味着类型检查只需在每个支持鸭子类型的函数内部发生一次,而使用 __array_function__ 时,每次调用 NumPy 函数都会发生类型检查。显然这取决于函数,但如果一个典型的支持鸭子数组的函数调用其他 NumPy 函数 3-5 次,那么这会增加 3-5 倍的开销。

  • 当使用 NumPy 数组 时,get_array_module 意味着每个函数额外多一次调用(__array_function__ 的开销保持不变),这意味着少量额外开销。

显式和隐式实现选择并非相互排斥的选项。实际上,我们熟悉的 NumPy API 覆盖的实现(通过 __array_function__,主要是 Dask、CuPy 和 Sparse,但不包括 Pint)也包含一种显式方式,通过直接导入模块来使用它们版本的 NumPy API(分别为 dask.arraycupysparse)。

局部、非局部与全局控制#

最终的设计轴心是用户如何控制 API 的选择

  • 局部控制,以多重调度和 Python 算术协议为例,通过检查类型或调用函数直接参数上的方法来确定使用哪个实现。

  • 非局部控制,例如 np.errstate,通过函数装饰器或上下文管理器使用全局状态来覆盖行为。控制是分层确定的,通过最内层的上下文。

  • 全局控制 提供了一种机制,供用户通过函数调用或配置文件设置默认行为。例如,matplotlib 允许全局选择绘图后端。

局部控制通常被认为是 API 设计的最佳实践,因为控制流是完全显式的,这使其最易于理解。非局部和全局控制偶尔也会使用,但通常是由于无知或缺乏更好的替代方案。

对于 NumPy 公共 API 的鸭子类型,我们认为非局部或全局控制将是错误的,主要是因为它们组合性不佳。如果一个库设置/需要一组覆盖,然后内部调用一个期望另一组覆盖的例程,结果行为可能会非常出乎意料。高阶函数尤其成问题,因为函数被求值的上下文可能与它们被定义的上下文不同。

我们认为非局部和全局控制适用的一个覆盖用例是选择一个保证具有完全一致接口的后端系统,例如 NumPy 数组上 numpy.fft 的更快替代实现。然而,这些超出了当前提案的范围,当前提案专注于鸭子数组。