NEP 18 — NumPy 高级数组函数的调度机制#

作者:

Stephan Hoyer <shoyer@google.com>

作者:

Matthew Rocklin <mrocklin@gmail.com>

作者:

Marten van Kerkwijk <mhvk@astro.utoronto.ca>

作者:

Hameer Abbasi <hameerabbasi@yahoo.com>

作者:

Eric Wieser <wieser.eric@gmail.com>

状态:

最终

类型:

标准轨道

创建:

2018-05-29

更新:

2019-05-25

决议:

https://mail.python.org/pipermail/numpy-discussion/2018-August/078493.html

摘要#

我们建议使用 __array_function__ 协议,允许 NumPy 函数的参数定义该函数如何操作它们。这将允许使用 NumPy 作为高效多维数组操作的高级 API,即使对于与 numpy.ndarray 大相径庭的数组实现也是如此。

详细描述#

NumPy 的高级 ndarray API 已在 NumPy 本身之外针对不同的架构多次实现,例如针对 GPU 数组 (CuPy)、稀疏数组 (scipy.sparse, pydata/sparse) 和并行数组 (Dask 数组),以及深度学习框架中各种类似 NumPy 的实现,例如 TensorFlow 和 PyTorch。

同样,许多项目都基于 NumPy API 构建,用于标记和索引数组 (XArray)、自动微分 (Autograd、Tangent)、掩码数组 (numpy.ma)、物理单位 (astropy.units、pint、unyt) 等,这些项目在 NumPy API 之上添加了附加功能。这些项目中的大多数也实现了 NumPy 高级 API 的接近变体。

我们希望能够将这些库一起使用,例如,我们希望能够将 CuPy 数组放在 XArray 中,或对 Dask 数组代码执行自动微分。如果为 NumPy ndarray 编写的代码也可以由其他类似 NumPy 的项目使用,则更容易实现这一点。

例如,我们希望以下代码示例能够与任何类似 NumPy 的数组对象一样好地工作

def f(x):
    y = np.tensordot(x, x.T)
    return np.mean(np.exp(y))

今天,NumPy 中的各种协议机制已经可以实现其中一部分。

  • np.exp 函数检查 __array_ufunc__ 协议

  • .T 方法使用 Python 的方法调度工作

  • np.mean 函数显式检查参数上的 .mean 方法

但是其他函数,例如 np.tensordot 不进行调度,而是可能会强制转换为 NumPy 数组(使用 __array__)协议,或直接出错。为了获得足够的 NumPy API 覆盖范围以支持 XArray 和 autograd 等下游项目,我们希望支持 NumPy 中的几乎所有函数,这需要比 __array_ufunc__ 更广泛的协议。我们希望有一个协议,允许 NumPy 函数的参数控制并将其执行转移到另一个函数(例如 GPU 或并行实现),这种方式在各个项目中都是安全且一致的。

实现#

我们建议在 NumPy 中添加对新协议 __array_function__ 的支持。

此协议旨在作为 NumPy 功能的总括,这些功能未被用于通用函数(如 np.exp)的 __array_ufunc__ 协议所涵盖。语义与 __array_ufunc__ 非常相似,只是操作由任意可调用对象而不是 ufunc 实例和方法指定。

原型实现可在 此笔记本 中找到。

警告

__array_function__ 协议及其在特定函数上的使用是实验性的。我们计划保留一个接口,使之可以覆盖 NumPy 函数,但是对特定函数进行此操作的方式**可能会并且将会改变**,而几乎不会发出警告。如果您不接受这种降低的后向兼容性保证,请不要依赖于非 NumPy 数组的 NumPy 函数重写。有关更多详细信息,请参见下面的“非目标”。

注意

使用 __array_function__ 协议进行调度已实现,但尚未默认启用

  • 在 NumPy 1.16 中,您需要在导入 NumPy 之前设置环境变量 NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=1 来测试 NumPy 函数重写。

  • 在 NumPy 1.17 中,该协议将默认启用,但可以使用 NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=0 禁用。

  • 最终,预计 __array_function__ 将始终启用。

接口#

我们建议对 __array_function__ 的实现使用以下签名

def __array_function__(self, func, types, args, kwargs)
  • func 是 NumPy 公共 API 公开的任意可调用对象,它以 func(*args, **kwargs) 的形式调用。

  • types 是来自原始 NumPy 函数调用的唯一参数类型的 集合,这些参数类型实现了 __array_function__

  • 元组 args 和字典 kwargs 直接从原始调用传递。

不同于__array_ufunc__,对于func的类型,或者argskwargs中哪些可能包含实现数组API的对象,没有高级别的保证。

为了方便__array_function__的实现者,types提供所有带有'__array_function__'属性的参数类型。这允许实现者快速识别他们应该将__array_function__的执行委托给其他参数的情况。types的类型故意含糊不清:frozenset最符合预期用途,但我们可能会出于性能原因使用tuple。无论如何,__array_function__的实现不应该依赖于types的迭代顺序,这将违反明确定义的“类型转换层次结构”(如NEP-13中所述)。

实现NumPy API的项目示例#

大多数__array_function__的实现将从两个检查开始

  1. 给定的函数是我们知道如何重载的吗?

  2. 所有参数都是我们知道如何处理的类型吗?

如果这些条件成立,__array_function__应该返回其针对func(*args, **kwargs)的实现调用的结果。否则,它应该返回哨兵值NotImplemented,表示这些类型未实现该函数。这比直接引发TypeError更好,因为它给了*其他*参数定义操作的机会。

__array_function__的返回值没有一般性要求,尽管大多数合理的实现可能应该返回与函数参数之一类型相同的数组。如果/当Python获得对协议的类型支持并且NumPy添加静态类型注释时,__array_function__@overload 实现将指示返回类型为Any

定义自定义装饰器(如下面的implements)来注册__array_function__实现也可能很方便。

HANDLED_FUNCTIONS = {}

class MyArray:
    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        # Note: this allows subclasses that don't override
        # __array_function__ to handle MyArray objects
        if not all(issubclass(t, MyArray) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)

def implements(numpy_function):
    """Register an __array_function__ implementation for MyArray objects."""
    def decorator(func):
        HANDLED_FUNCTIONS[numpy_function] = func
        return func
    return decorator

@implements(np.concatenate)
def concatenate(arrays, axis=0, out=None):
    ...  # implementation of concatenate for MyArray objects

@implements(np.broadcast_to)
def broadcast_to(array, shape):
    ...  # implementation of broadcast_to for MyArray objects

请注意,__array_function__实现不需要包含对应NumPy函数的所有可选参数(例如,上面的broadcast_to省略了无关的subok参数)。只有在NumPy函数调用中显式使用了可选参数,才会将其传递给__array_function__

注意

就像__add__之类的内置特殊方法一样,编写正确的__array_function__方法应该始终在遇到未知类型时返回NotImplemented。否则,如果操作还包含你的对象之一,将无法正确地从另一个对象重写NumPy函数。

NumPy代码库本身的必要更改#

这需要在NumPy代码库中进行两处更改

  1. 一个函数来检查可用的输入,查找这些输入上的__array_function__属性,并适当地调用这些方法,直到一个方法成功。这在所有都是NumPy数组的常见情况下需要很快,即使重载输入的数量很大(例如,对于np.concatenate可能就是这种情况),也需要具有可接受的性能(不低于线性时间)。

    这是一个中等复杂度的附加函数。

  2. 在所有相关的NumPy函数中调用此函数。

    这影响NumPy代码库的许多部分,但复杂性非常低。

查找和调用正确的__array_function__#

给定一个NumPy函数、*args**kwargs输入,我们需要在*args**kwargs中搜索所有可能具有__array_function__属性的适当输入。然后,我们需要在这些可能的方法中进行选择并执行正确的方法。在几个可能的实现之间进行协商可能很复杂。

查找参数#

有效参数可以直接位于*args**kwargs中,例如np.tensordot(left, right, out=out)的情况,或者它们可能嵌套在列表或字典中,例如np.concatenate([x, y, z])的情况。这可能有两个问题:

  1. 一些函数被赋予很长的值列表,遍历它们可能非常昂贵。

  2. 一些函数可能有一些我们不想检查的参数,即使它们具有__array_function__方法。

为了解决这些问题,NumPy函数应该明确指示哪些参数可以被重载,以及如何检查这些参数。通常,这应该包括所有记录为array_likendarray的参数。

我们建议通过为每个重载的NumPy函数编写“调度器”函数来做到这一点

  • 这些函数将使用传递给NumPy函数的完全相同的参数调用(即dispatcher(*args, **kwargs)),并且应该返回要检查重载的参数的迭代器。

  • 调度器函数必须与其对应的NumPy函数具有完全相同的 positional、optional 和 keyword-only 参数。否则,调用NumPy函数的有效调用可能会在调用其调度器时导致错误。

  • 因为关键字参数的默认*值*没有__array_function__属性,按照约定,我们将所有默认参数值设置为None。这降低了签名不同步的可能性,并最大限度地减少了调度器中的额外信息。唯一的例外情况应该是参数值以某种方式影响调度的案例,这种情况应该很少见。

np.concatenate调度器的示例可能具有指导意义

def _concatenate_dispatcher(arrays, axis=None, out=None):
    for array in arrays:
        yield array
    if out is not None:
        yield out

concatenate 调度器被编写为生成器函数,这允许它潜在地包含可选out参数的值,而无需使用要连接的对象的(可能很长)列表创建新的序列。

尝试__array_function__方法直到找到正确的方法#

许多参数可以实现__array_function__协议。其中一些可能决定,鉴于可用的输入,它们无法确定正确的结果。我们如何调用正确的方法?如果几个都是有效的,那么哪个优先?

在大多数情况下,使用__array_function__进行调度的规则与使用__array_ufunc__的规则相匹配(参见NEP-13)。特别是

  • NumPy将从所有指定的输入中收集__array_function__的实现,并按顺序调用它们:子类优先于超类,否则从左到右。请注意,在涉及子类的一些极端情况下,这与Python的当前行为略有不同。

  • __array_function__的实现通过返回NotImplemented以外的任何值来表明它们可以处理该操作。

  • 如果所有__array_function__方法都返回NotImplemented,NumPy将引发TypeError

如果没有__array_function__方法,NumPy将默认调用其自身实现,该实现旨在用于NumPy数组。例如,当所有类似数组的参数都是Python数字或列表时,就会出现这种情况。(NumPy数组确实有一个__array_function__方法,如下所示,但如果任何参数(而不是NumPy数组子类)实现了__array_function__,它将始终返回NotImplemented。)

与当前__array_ufunc__的行为相比,一个偏差是NumPy将只调用每个唯一类型的_第一个_参数的__array_function__。这与Python的调用反射方法的规则相匹配,并且这确保了即使有很多重载参数,检查重载也能具有可接受的性能。为了避免这两个调度协议之间长期存在差异,我们应该更新__array_ufunc__以匹配此行为。

numpy.ndarray上的__array_function__方法#

具有__array_function__的子类的用例与具有__array_ufunc__的子类的用例相同,因此numpy.ndarray也定义了一个__array_function__方法。

def __array_function__(self, func, types, args, kwargs):
    if not all(issubclass(t, ndarray) for t in types):
        # Defer to any non-subclasses that implement __array_function__
        return NotImplemented

    # Use NumPy's private implementation without __array_function__
    # dispatching
    return func._implementation(*args, **kwargs)

此方法与NumPy的调度规则相匹配,因此在大多数情况下,可以假装ndarray.__array_function__不存在。下面在array_function_dispatch装饰器中定义的私有_implementation属性允许我们避免在__array_ufunc__协议中需要的NumPy数组的特殊情况。

__array_function__协议总是先调用子类,然后再调用超类,因此如果任何ndarray子类参与运算,它们将有机会覆盖它,就像任何其他参数覆盖__array_function__一样。但是,结合基本NumPy数组和子类的操作中的默认行为是不同的:如果子类返回NotImplemented,则将调用NumPy对该函数的实现,而不是引发异常。这是合适的,因为子类应该可以替换

我们仍然告诫子类的作者,在依赖NumPy内部实现的细节时要谨慎。并非总是可以编写一个完全可替换的ndarray子类,例如,在涉及创建新数组的情况下,尤其是因为NumPy使用了专门针对基本NumPy数组的内部优化,例如用C编写的代码。即使NumPy的实现今天恰好有效,将来也可能无效。在这些情况下,你的补救措施是通过你子类上的__array_function__重新实现顶级NumPy函数。

NumPy函数内的更改#

给定一个定义上述行为的函数,现在将其称为implement_array_function,我们现在需要从每个相关的NumPy函数中调用该函数。这是一个普遍的更改,但代码相当简单且无害,如果没有任何参数实现__array_function__协议,则应快速完成且不会产生影响。

为了实现这一点,我们定义了一个array_function_dispatch装饰器来重写NumPy函数。基本实现如下:

def array_function_dispatch(dispatcher, module=None):
    """Wrap a function for dispatch with the __array_function__ protocol."""
    def decorator(implementation):
        @functools.wraps(implementation)
        def public_api(*args, **kwargs):
            relevant_args = dispatcher(*args, **kwargs)
            return implement_array_function(
                implementation, public_api, relevant_args, args, kwargs)
        if module is not None:
            public_api.__module__ = module
        # for ndarray.__array_function__
        public_api._implementation = implementation
        return public_api
    return decorator

# example usage
def _broadcast_to_dispatcher(array, shape, subok=None):
    return (array,)

@array_function_dispatch(_broadcast_to_dispatcher, module='numpy')
def broadcast_to(array, shape, subok=False):
    ...  # existing definition of np.broadcast_to

使用装饰器很棒!我们不需要更改现有NumPy函数的定义,只需要为调度程序函数编写几行附加代码。我们甚至可以对具有相同签名的函数族(例如sumprod)重用单个调度程序。对于此类函数,最大的更改可能是向文档字符串添加几行,以说明检查哪些参数的重载。

特别值得一提的是装饰器对functools.wraps的使用。

  • 这确保了包装后的函数具有与包装后的NumPy函数相同的名称和文档字符串。

  • 在Python 3上,它还确保装饰器函数复制原始函数签名,这对于基于自省的工具(如自动完成)非常重要。

  • 最后,它确保包装后的函数可以被pickle

示例用法说明了为NumPy贡献者编写调度程序的几种最佳实践。

  • 我们传递了module参数,该参数反过来设置生成的函数上的__module__属性。这是为了更好地显示错误消息,这里用于NumPy在找不到实现时内部引发的错误,例如TypeError: no implementation found for 'numpy.broadcast_to'。将__module__设置为NumPy公共API中的规范位置鼓励用户使用NumPy的公共API在__array_function__中标识函数。

  • 调度程序是一个返回元组的函数,而不是使用yield的等效(同样有效)的生成器。

    # example usage
    def broadcast_to(array, shape, subok=None):
        yield array
    

    这并非偶然:当调度程序函数返回内置序列类型(tuplelist)时,NumPy对__array_function__的调度实现速度最快。

    相关的是,即使在某些情况下你_知道_它们不可能具有__array_function__方法,调度程序返回参数也是完全可以的。这可能发生在具有默认参数(例如None)或复杂签名的函数中。NumPy的调度逻辑非常快速地解决了这些情况,因此通常不值得自己费力地解析它们。

注意

上面array_function_dispatch的代码已从该NEP的原始版本更新,以匹配NumPy中的实际实现

可扩展性#

这种方法的一个重要优点是它允许向NumPy函数添加新的可选参数,而不会破坏已经依赖__array_function__的代码。

这不是一个理论上的问题。NumPy在函数内部(如np.sum())中较旧、杂乱的覆盖实现,在我们决定添加新的可选参数时,需要一些笨拙的技巧,例如,新的keepdims参数仅在使用时才传递。

def sum(array, ..., keepdims=np._NoValue):
    kwargs = {}
    if keepdims is not np._NoValue:
        kwargs['keepdims'] = keepdims
    return array.sum(..., **kwargs)

对于__array_function__实现者来说,这也意味着可以逐步实现甚至现有的可选参数,并且仅在有意义的情况下实现。例如,实现不可变数组的库不需要在函数签名中显式包含不受支持的out参数。这在实现上可能有些繁琐,例如:

def my_sum(array, ..., out=None):
    if out is not None:
        raise TypeError('out argument is not supported')
    ...

因此,我们避免鼓励向NumPy调用的函数签名添加万能的**ignored_kwargs的诱人捷径,因为这会对拼写错误或忽略的参数静默失败。

性能#

性能始终是NumPy的一个关注点,即使NumPy用户已经通过选择Python语言本身将可用性优先于纯粹的速度。重要的是,这个新的__array_function__协议在NumPy函数作用于NumPy数组的典型情况下不会造成很大的成本。

我们的微基准测试结果表明,上面描述的覆盖机制的纯Python实现为每个NumPy函数调用增加了大约2-3微秒的开销,而没有任何重载参数。作为参考,小型数组上的典型NumPy函数的运行时间为1-10微秒,主要取决于函数逻辑中有多少部分是用C编写的。例如,1微秒大约是ndarray.sum()方法(1.6 us)和numpy.sum()函数(2.6 us)之间的速度差异。

幸运的是,我们预计implement_array_function的C实现的开销会小得多,而大部分运行时间都在这里。这将使array_function_dispatch装饰器和调度程序函数本身增加大约0.5微秒的开销,在典型情况下可能约为1微秒的开销。

在我们看来,对于用Python编写的代码来说,接受这种级别的开销是合理的。我们非常确定,绝大多数NumPy用户并不关心用微秒来衡量的NumPy函数的性能差异,因为用Python在不到1微秒的时间内完成_任何_事情都很困难。

NumPy之外的使用#

此协议中没有任何内容是NumPy特有的。我们是否应该鼓励第三方库使用相同的__array_function__协议来重载非NumPy函数,例如,用于在SciPy中实现数组实现通用功能?

这将提供显著的优势(SciPy 不需要发明自己的分发系统),并且我们想不到任何缺点,因为每个使用 __array_function__ 进行分发的函数都需要被明确识别。像 Dask、CuPy 和 Autograd 这样的库已经类似于它们包装 NumPy 的方式,包装了 SciPy 功能的有限子集(例如,scipy.linalg)。

如果我们想这样做,我们应该至少公开装饰器 array_function_dispatch(),并可能也公开更低级别的 implement_array_function() 作为 NumPy 公共 API 的一部分。

非目标#

我们的目标是制定一个基本的策略,该策略可以相对机械地在 NumPy API 中几乎所有函数上应用,并在相对较短的时间内(单个 NumPy 版本的开发周期)完成。

我们希望在第一次尝试中同时获得 __array_function__ 协议和所有特定的重载,但我们的明确目标是获得一个大部分有效的方案(并且可以迭代改进),而不是等待最佳实现。快速行动的代价是,目前 **此协议应被视为严格的实验性**。我们保留随时更改此协议的细节以及将来特定 NumPy 函数如何使用它的权利——即使是在仅修复错误的 NumPy 版本中也是如此。实际上,一旦解决了 __array_function__ 的初始问题,我们将使用缩短的弃用周期,短至单个主要 NumPy 版本(例如,短至四个月)。

特别是,我们不打算编写额外的 NEP 来列出所有需要重载的特定函数,以及它们应该如何重载。我们将把这留给各个拉取请求的提交者的自由裁量权,相信他们会将任何争议提交给相关方讨论。

但是,我们已经知道几类函数应该明确地从 __array_function__ 中排除。这些将需要它们自己的协议。

  • 通用函数,它们已经有自己的协议。

  • arrayasarray,因为它们明确用于强制转换为实际的 numpy.ndarray 对象。

  • 任何类型的函数方法的分发,例如,np.random.RandomState 对象的方法。

我们还预计,最初将使用 __array_function__ 协议来覆盖特定函数的机制将来可能会发生变化。作为我们预期将来如何破坏行为的一个具体示例,一些函数(例如 np.where)目前不是 NumPy 通用函数,但在将来可能会成为通用函数。当/如果发生这种情况时,我们将更改此类重载,使其从使用 __array_function__ 变为使用更专业的 __array_ufunc__

向后兼容性#

此提案不会更改现有语义,除了那些当前具有 __array_function__ 属性的参数,这种情况应该很少见。

替代方案#

专用协议#

我们可以(也应该)继续开发像 __array_ufunc__ 这样的协议,用于 NumPy 功能的内聚子集。

如上所述,如果这意味着我们使用 __array_function__ 重载的某些函数应该改为使用新的协议,只要 __array_function__ 保持其实验状态,这都是明确允许的。

切换到新的协议应该使用 NumPy 正常弃用周期的缩短版本。

  • 对于单个主要版本,在检查任何新协议后,NumPy 仍然应该检查实现给定函数的 __array_function__ 方法。如果任何参数从 __array_function__ 返回的值不是 NotImplemented,则应发出描述性的 FutureWarning

  • 在下一个主要版本中,将删除对 __array_function__ 的检查。

单独的命名空间#

重载函数的单独命名空间是另一种可能性,无论是在 NumPy 内部还是外部。

这具有减轻任何关于向后兼容性的潜在问题的优势,并将为快速实验提供最大的自由。从长远来看,它将提供一个干净的抽象层,将 NumPy 的高级 API 与 numpy.ndarray 对象上的默认实现分开。

缺点是这将需要所有现有代码的显式选择加入,例如 import numpy.api as np,并且从长远来看会导致维护两个单独的 NumPy API。此外,numpy 本身中的许多函数已经被重载(但不足),因此关于 NumPy 中高级和低级 API 的混淆仍然存在。

或者,可以为 NumPy 高级 API 的非重载版本创建一个单独的命名空间,例如 numpy.array_only,用于 NumPy 数组性能至关重要的场合。这与单独命名空间的大部分缺点相同。

多重分发#

对我们建议的 __array_function__ 协议的替代方案是将 NumPy 的核心函数实现为 多方法。虽然我们中的一位为 Python 编写了一个 多重分发库,但我们认为这种方法在短期内对 NumPy 没有意义。

主要原因是 NumPy 已经有一个经过验证的分发机制 __array_ufunc__,它基于 Python 自身用于算术运算的分发系统,添加另一个以非常不同方式工作的机制会令人困惑。这也会对 NumPy 本身造成更具侵入性的更改,NumPy 需要获得多重分发实现。

NumPy 高级 API 的多重分发实现将来可能会有意义。幸运的是,__array_function__ 并不会排除这种可能性,因为可以用多重分发来编写默认 __array_function__ 实现的垫片。

基于有限核心 API 的实现#

一些 NumPy 函数的内部实现非常简单。例如

  • np.stack() 通过将索引与 np.newaxisnp.concatenateshape 属性组合,仅用几行代码实现。

  • np.mean() 在内部使用 np.sum()np.divide().astype().shape 实现。

这表明可以定义一个最小的“核心”ndarray 接口,并在 NumPy 内部依赖它来实现完整的 API。这是一个有吸引力的选择,因为它可以显著减少新数组实现所需的工作量。

但是,这也有一些缺点。

  1. NumPy 如何使用重载函数实现高级函数的细节现在成为 NumPy 公共 API 的隐式部分。例如,将 stack 重构为在内部使用 np.block() 而不是 np.concatenate() 将成为一个破坏性更改。

  2. 数组库可能更喜欢以不同于 NumPy 的方式实现高级函数。例如,库可能更喜欢直接实现像 mean() 这样的基本运算,而不是依赖于先进行求和然后除法。更一般地说,目前尚不清楚什么确切地构成核心功能,弄清楚这一点可能是一个大型项目。

  3. 我们还没有用于数组对象的属性和方法重载的系统,例如,用于访问 .dtype.shape。这应该是未来 NEP 的主题,但在那之前,我们应该谨慎依赖这些属性。

考虑到这些问题,我们认为支持显式重载 NumPy API 中几乎所有公共函数是有价值的。这并不排除将来使用 __array_function__ 和协议和/或基类来重写 NumPy 函数以简化核心功能,以确保数组公开像 numpy.ndarray 一样的方法和属性。但是,要使其正常工作,这需要能够使用 __array_function__ 实现 *一些* 但不是所有函数的可能性,例如下一节所述。

NumPy API 的部分实现#

在当前设计中,实现 __array_function__ 以重载至少一个函数的类隐式声明了实现整个 NumPy API 的意图。不可能只在一个类型上实现 np.concatenate(),而对于所有其他函数则回退到使用 np.asarray() 进行类型转换的 NumPy 默认行为。

这可能会带来向后兼容性问题,从而阻止库以增量方式采用__array_function__。例如,目前大多数 NumPy 函数会隐式地将pandas.Series对象转换为 NumPy 数组,许多 pandas 用户肯定依赖于这种行为。如果 pandas 只为np.concatenate实现了__array_function__,那么像np.nanmean这样的无关 NumPy 函数会突然在 pandas 对象上因引发 TypeError 而中断。

即使是重新实现了大部分 NumPy 公共 API 的库,有时也依赖于使用 NumPy 的实用函数而无需包装器。例如,CuPy 和 JAX 都只是使用了别名np.result_type,后者已经通过dtype属性支持鸭子类型。

使用__array_ufunc__,可以通过将所有参数转换为 NumPy 数组并重新调用 ufunc 来缓解此问题,但是__array_function__支持的异构函数签名使得不可能为__array_function__实现这种通用的回退行为。

我们考虑了三种可能的解决方法,但没有一种完全令人满意。

  1. 更改所有返回NotImplemented的参数的含义,从__array_function__表示应将所有参数强制转换为 NumPy 数组并重试操作。但是,许多数组库(例如,scipy.sparse)真的不希望隐式转换为 NumPy 数组,并且经常为了这个原因而避免实现__array__。隐式转换可能导致无声错误和性能下降。

    潜在地,我们只可以对实现了__array__的类型启用此行为,这将解决像 scipy.sparse 这样最棘手的情况。但在实践中,很大一部分提供像 NumPy 数组一样的高级 API 的类已经实现了__array__。这将阻止在这些对象上可靠地使用 NumPy 的高级 API。

  2. 使用某种其他哨兵值,例如np.NotImplementedButCoercible,来指示实现 NumPy 部分高级数组 API 的类可以作为回退进行强制转换。如果所有参数都返回NotImplementedButCoercible,则将强制转换参数并重试操作。

    不幸的是,遇到NotImplementedButCoercible后的正确行为并不总是显而易见的。特别具有挑战性的是“混合”情况,其中一些参数返回NotImplementedButCoercible,而其他参数返回NotImplemented。只强制转换“可强制转换”的参数后是否会重试分派?如果是这样,那么我们可以想象会任意多次循环遍历分派逻辑。无论哪种方式,分派规则肯定会变得更加复杂,更难以理解。

  3. 允许访问 NumPy 函数的实现,例如,在 NumPy 函数上以公开的__skip_array_function__属性的形式。这允许通过在__array_function__方法内使用func.__skip_array_function__来回退到 NumPy 的实现,并且还可以潜在地用于避免分派的开销。但是,它存在可能公开 NumPy 函数实现细节的风险,而这些 NumPy 函数在内部并没有调用np.asarray()。请参阅此说明,了解完整讨论的总结。

这些解决方案可以解决实际用例,但代价是增加了复杂性。在做出难以回滚的决定之前,我们希望获得关于__array_function__实际使用方式的经验。

检查类型注释的魔法装饰器#

原则上,Python 3 类型注释包含足够的信息来自动创建大多数dispatcher函数。使用这些注释来避免手动编写分派器会很方便,例如:

@array_function_dispatch
def broadcast_to(array: ArrayLike
                 shape: Tuple[int, ...],
                 subok: bool = False):
    ...  # existing definition of np.broadcast_to

这需要某种形式的自动代码生成,无论是在编译时还是导入时。

我们认为这是一个将来可以考虑的有趣的扩展。我们认为现在这样做没有意义,因为代码生成涉及权衡,并且 NumPy 使用类型注释的经验仍然非常有限。即使 NumPy 只是 Python 3(这将在2019 年的某个时候发生),我们现在还没有准备好直接注释 NumPy 的代码库。

对实现特定参数的支持#

我们可以允许__array_function__实现通过在分派器函数中包含**ignored_kwargs来添加它们自己的可选关键字参数,例如:

def _concatenate_dispatcher(arrays, axis=None, out=None, **ignored_kwargs):
    ...  # same implementation of _concatenate_dispatcher as above

实现特定参数在其他方面模拟 NumPy 高级 API 的库中比较常见(例如,dask.array.sum()添加了split_every,而tensorflow.reduce_sum()添加了name)。在 NumPy 中支持它们对于在 NumPy 函数之上实现新的高级数组函数的库特别有用,例如:

def mean_squared_error(x, y, **kwargs):
    return np.mean((x - y) ** 2, **kwargs)

否则,我们需要为每个数组实现分别创建mean_squared_error的单独版本,以便将实现特定参数传递给mean()

我们不允许添加可选的位置参数,因为这些参数保留供 NumPy 本身将来使用,但关键字参数之间的冲突应该相对较少。

但是,这种灵活性是有代价的。特别是,它隐式地向所有包装的 NumPy 函数的签名添加了**kwargs,而实际上并没有包含它(因为我们使用functools.wraps)。这意味着它不太可能与静态分析工具一起良好工作,这些工具可能会报告无效参数。同样,可读性也有一定的代价:这些可选参数不会包含在 NumPy 函数的文档字符串中。

目前还不清楚这种权衡是否值得,因此我们建议暂时将其排除在外。添加实现特定参数将需要直接使用那些库。

协议的其他可能选择#

数组函数__array_function__只包含两个参数functypes,它们提供有关函数调用上下文的信息。

func是协议的一部分,因为无法避免它:实现需要能够通过将函数与 NumPy 的公共 API 匹配来进行分派。

types包含在内,因为我们几乎可以免费地将其作为收集__array_function__实现的一部分来计算,以在implement_array_function中调用。我们还认为它将被许多__array_function__方法使用,否则这些方法需要自己提取这些信息。提供每种类型的单个实例同样容易,但是只提供类型似乎更简洁。

更进一步,有人建议__array_function__应该是一个classmethod。我们同意删除冗余的self参数会更简洁一些,但认为这种微小的改进不值得破坏__array_ufunc__的优先级。

我们认为还有另外两个参数可能重要地传递给__array_ufunc__实现。

  • ndarray.__array_function__中访问非分派实现(即在用array_function_dispatch包装之前)将允许我们从implement_array_function中删除该方法的特殊情况逻辑。

  • 访问传递到array_function_dispatch()中的dispatcher函数将允许__array_function__实现通过调用dispatcher(*args, **kwargs)以通用的方式确定“类数组”参数的列表。这对于基于数组属性(例如dtypeunits)的值而不是直接基于数组类型进行分派的__array_function__实现可能很有用。

我们现在将其排除在外,因为我们不知道它们是否必要。如果我们想将来包含它们,最简单的方法是更新array_function_dispatch装饰器以将其作为函数属性添加。

运行时生成的 Callable 对象#

NumPy 有一些 API 可以动态定义可调用对象,例如 vectorizerandom.RandomState 对象的方法。在科学 Python 栈的其他核心库中也可以找到类似的例子,例如 scipy.stats 中的分布对象和 scikit-learn 中的模型对象。能够为这些可调用对象编写重载函数也很不错。但这对 __array_function__ 协议提出了挑战,因为与函数不同,在 numpy 命名空间中没有公共对象可以传递到 func 参数。

我们可以通过建立一个关于如何检查 func 参数的替代约定来解决这个问题,例如,使用 func.__self__ 获取类对象,并使用 func.__func__ 返回未绑定的函数对象。但是,需要注意的是,这会将当前的实现细节作为接口的永久特性,例如 vectorize 是作为类而不是闭包实现的,或者方法是直接实现还是使用描述符实现。

鉴于其复杂性和有限的用例,我们目前也推迟处理这个问题,但我们相信,如果需要的话,__array_function__ 未来可以扩展以适应这些用例。

讨论#

在一些 GitHub issues 中讨论了该提案的各种替代方案

  1. pydata/sparse #1

  2. numpy/numpy #11129

此外,它也是一篇博文的主题。在此之后,它在NumPy 开发者冲刺(在加州大学伯克利分校数据科学研究所 (BIDS)举行)中进行了讨论。

该提案本身的详细讨论可以在邮件列表和相关的 pull requests 上找到(123