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/[email protected]/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__，部分原因是担心破坏现有代码：用户期望 NumPy 函数（如 np.concatenate）返回 NumPy 数组。这是 __array_function__ 设计的一个基本限制，我们选择允许覆盖现有的 numpy 命名空间。像 Dask 和 CuPy 这样的库已经查看并接受了 __array_function__ 带来的向后兼容性影响；如果这种影响不存在，对它们来说仍然会更好。

      请注意，像 PyTorch 和 scipy.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.arange 和 np.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_module。get_array_module 执行与 __array_function__ 相同的类型解析，并返回一个模块，该模块的 API 保证与 numpy 的标准接口匹配，可以实现对所有提供的数组类型的操作。

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

新协议是选择加入的、明确的并且具有本地控制；有关这些设计特性的重要性，请参阅 附录：API 重写的设计选择。

数组模块契约#

由 get_array_module/__array_module__ 返回的模块应尽最大努力在新的数组类型(s) 上实现 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__ 属性的参数。

特殊方法应返回一个具有与 numpy 匹配的 API 的命名空间，或者返回 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 ufunc（例如，来自 Numba 或 SciPy）仍然很有用。

鉴于它们现有的采用和不同的用例，我们认为目前没有必要删除或弃用 __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)

为了更容易编写鸭子数组，我们也可以将这些 Mixin 类添加到 numpy.lib.mixins 中（但上面的示例可能就足够了）。

考虑的替代方案#

命名#

我们喜欢 __array_module__ 这个名称，因为它反映了现有的 __array_function__ 和 __array_ufunc__ 协议。另一个合理的选项可能是 __array_namespace__。

调用此协议的 NumPy 函数应该叫什么（在本提案中为 get_array_module）不太清楚。一些可能的替代方案：array_module、common_array_module、resolve_array_module、get_namespace、get_numpy、get_numpylike_module、get_duck_array_module。

请求 NumPy API 的受限子集#

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

我们认为定义 NumPy API 的“最小”子集（或子集）将是有益的练习，省略很少使用或不推荐的功能。例如，最小 NumPy 可能包含 stack，但不包含其他堆叠函数 column_stack、dstack、hstack 和 vstack。这可以清楚地向鸭子数组作者和用户表明哪些功能是核心功能，哪些功能可以跳过。

支持请求 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 进行测试并在任何有效的鸭子数组库中使用，当仅对 NumPy 数组调用 get_array_module 时，NumPy 本身将返回 numpy 模块的受限版本。省略的函数将根本不存在。

不幸的是，我们还没有确定这些受限子集应该是什么，因此现在这样做没有意义。当/如果我们这样做时，我们可以向 get_array_module 添加新的关键字参数，或者添加新的顶级函数，例如 get_minimal_array_module。我们还需要添加一个新的基于 __array_module__ 的协议（例如，__array_module_minimal__），或者可以向 __array_module__ 添加一个可选的第二个参数（使用 try/except 捕获错误）。

用于隐式分派的新命名空间#

与其使用 __array_function__ 在主 numpy 命名空间中支持覆盖，我们可以创建一个新的选择加入命名空间，例如 numpy.api，其中包含支持分派的 NumPy 函数版本。这些覆盖需要新的选择加入协议，例如，基于 __array_function__ 的 __array_function_api__。

这将通过选择加入来解决 __array_function__ 最大的限制，并且还允许明确覆盖像 asarray 这样的函数，因为 np.api.asarray 将始终表示“转换类数组对象”。但它不会解决 __array_module__ 满足的所有分派需求，并且将使我们不得不为数组用户和实现者支持一个相当复杂的协议。

我们可能可以通过 __array_module__ 协议实现这样一个新的命名空间。当然，一些用户会发现这很方便，因为它稍微减少了样板代码。但这会让用户面临一个令人困惑的选择：何时使用 get_array_module 与 np.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__ 的开销保持不变），这意味着少量的额外开销。

实现的显式选择和隐式选择不是相互排斥的选项。事实上，我们熟悉的大多数通过 __array_function__ 覆盖 NumPy API 的实现（即 Dask、CuPy 和 Sparse，但不包括 Pint）也包括一种显式的方法来通过直接导入模块来使用它们自己的 NumPy API 版本（分别为 dask.array、cupy 或 sparse）。

本地控制、非本地控制和全局控制#

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

  • **本地控制**，如多重分派和 Python 算术协议所示，通过检查类型或在函数的直接参数上调用方法来确定使用哪个实现。

  • **非本地控制**，例如 np.errstate，通过函数装饰器或上下文管理器以全局状态覆盖行为。控制是按层次结构确定的，通过最内部的上下文。

  • **全局控制**为用户提供了一种机制来设置默认行为，无论是通过函数调用还是配置文件。例如，matplotlib 允许设置绘图后端的全局选择。

本地控制通常被认为是 API 设计的最佳实践，因为控制流是完全明确的，这使得它最容易理解。偶尔会使用非本地和全局控制，但通常是由于无知或缺乏更好的替代方案。

在 NumPy 公共 API 的鸭子类型的情况下，我们认为非本地或全局控制将是错误的，主要是因为它们**无法很好地组合**。如果一个库设置/需要一组覆盖，然后在内部调用一个期望另一组覆盖的例程，则结果行为可能会非常令人惊讶。高阶函数尤其成问题，因为函数求值所在的上下文可能不是函数定义所在的上下文。

我们认为，在某些覆盖用例中，非本地和全局控制是合适的，例如选择一个保证具有完全一致接口的后端系统，例如 numpy.fft 在 NumPy 数组上的更快替代实现。但是，这些超出了当前提案的范围，当前提案重点关注鸭子数组。

本页内容