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__ 对向后不兼容的影响；如果这种影响不存在，对它们来说会更好。

      请注意，像 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__ 相同的类型解析，并返回一个具有承诺匹配 numpy 标准接口的 API 的模块，该模块可以实现所有提供的数组类型的操作。

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

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

数组模块约定#

由 get_array_module/__array_module__ 返回的模块应尽最大努力在新数组类型上实现 NumPy 的核心功能。未实现的功能应直接省略（例如，访问未实现的功能应引发 AttributeError）。未来，我们预计将制定一个请求 NumPy API 受限子集的协议；有关更多详细信息，请参阅请求 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__，它是一个 Python 集合，包含传递给 get_array_module 的所有具有 __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__ 的混合类#

尽管用户界面存在差异，__array_module__ 和实现 NumPy API 的模块仍然包含实现现有“鸭子数组”协议分派所需的功能。

例如，以下混合类将为这些特殊方法提供合理的默认值，以 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_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 捕获错误）。

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

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

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

我们**可能**可以通过 __array_module__ 协议来实现这样一个新的命名空间。当然，一些用户会觉得这很方便，因为它稍微少了一些样板代码。但这会让用户面临一个令人困惑的选择：何时使用 get_array_module vs np.api.something。此外，我们还需要添加和维护一个全新的模块，这比仅仅添加一个函数要昂贵得多。

基于类型和数组进行分派，而不是仅基于类型#

除了仅通过唯一的数组类型支持分派外，我们还可以通过数组对象支持分派，例如，通过将 arrays 参数作为 __array_module__ 协议的一部分传递。这可能有助于 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__ 的三个主要设计决策轴。

用户选择加入 vs. 选择退出#

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

理论上，这种方法降低了用户代码和库采用这些协议的门槛，因为使用标准 NumPy 命名空间的代码是自动兼容的。但实际上，这并没有奏效。例如，大多数维护良好的使用 NumPy 的库都遵循使用 np.asarray() 转换所有输入的最佳实践，而他们不得不明确放宽该限制才能使用 __array_function__。我们的经验是，使库与新的“鸭子数组”类型兼容通常需要至少少量的工作来适应数据模型和可以有效实现的运算的差异。

这些选择退出的方法也大大增加了采用这些协议的库的向后兼容性复杂性，因为作为库选择加入，它们也选择了用户，无论用户是否期望。为了赢得那些无法采用 __array_function__ 的库，选择加入的方法似乎是必须的。

显式 vs. 隐式选择实现#

__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）也通过直接导入模块（分别是 dask.array、cupy 或 sparse）包含了一种显式使用它们 NumPy API 版本的方式。

本地 vs. 非本地 vs. 全局控制#

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

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

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

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

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

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

我们认为，在选择一个保证拥有完全一致接口的后端系统时，非局部和全局控制是合适的，例如 NumPy 数组上 numpy.fft 的更快的替代实现。然而，这些超出了当前提案的范围，当前提案的重点是“duck arrays”（鸭子数组）。