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 array),以及深度学习框架中各种类似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__ 的支持。

该协议旨在作为通用函数(如 np.exp)的 __array_ufunc__ 协议未涵盖的NumPy功能的包罗万象的机制。其语义与 __array_ufunc__ 非常相似,只是操作由任意可调用对象指定,而不是ufunc实例和方法。

原型实现可以在此notebook中找到。

警告

__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添加静态类型注解时,@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函数共享完全相同的位置参数、可选参数和仅关键字参数。否则,对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微秒)与 numpy.sum() 函数(2.6微秒)之间的速度差异。

幸运的是,我们预计使用 implement_array_function 的C实现会显著减少开销,因为大部分运行时间都集中在那里。这将使得 array_function_dispatch 装饰器和调度器函数本身增加大约0.5微秒的开销,在典型情况下可能总共增加约1微秒的开销。

在我们看来,对于用Python编写的代码,这种级别的开销是合理的。我们很确定绝大多数NumPy用户不关心NumPy函数上以微秒计的性能差异,因为在Python中做任何事都很难在不到一微秒内完成。

在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发布周期内)相对机械地应用于NumPy API中的几乎所有函数。

我们希望 __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__ 属性的参数(这种情况应该很少见)之外。

替代方案#

专用协议#

我们应该(并且也能够)继续为NumPy功能的内聚子集开发诸如 __array_ufunc__ 之类的协议。

如上所述,如果这意味着我们使用 __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与低级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() 这样的基本操作,而不是依赖于 sum() 后再进行除法。更普遍地说,目前尚不清楚究竟什么才算核心功能,弄清这一点可能是一个大项目。

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

鉴于这些担忧,我们认为支持NumPy API中几乎所有公共函数的显式重载是很有价值的。这并不排除未来将NumPy函数重写为具有 __array_function__ 的简化核心功能,以及一个协议和/或基类,以确保数组暴露类似 numpy.ndarray 的方法和属性。然而,为了良好运行,这将需要能够使用 __array_function__ 实现部分而非全部函数,例如,如下一节所述。

NumPy API的部分实现#

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

这可能会带来向后兼容性问题,从而阻碍库以增量方式采用 __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. 更改所有从 __array_function__ 返回 NotImplemented 的参数的含义,以表明所有参数都应强制转换为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函数内部实现细节的风险,对于那些内部不调用 np.asarray() 的NumPy函数而言。有关完整讨论的摘要,请参阅此说明

这些解决方案可以解决实际用例,但代价是增加了额外的复杂性。我们希望在做出难以回滚的决定之前,先了解 __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_everytensorflow.reduce_sum() 添加了 name)。在NumPy中支持它们对于在NumPy函数之上实现新的高级数组函数的库将特别有用,例如:

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

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

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

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

尚不清楚这种权衡是否值得,因此我们建议暂时不予考虑。添加特定于实现的参数将需要直接使用那些库。

协议的其他可能选择#

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

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

types 被包含在内,因为在 implement_array_function 中收集要调用的 __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 装饰器,将它们作为函数属性添加。

运行时生成的可调用对象#

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

我们或许可以通过建立另一种约定来处理这个问题,即如何检查 func 参数,例如,使用 func.__self__ 获取类对象,使用 func.__func__ 返回未绑定的函数对象。然而,需要谨慎,因为这会将目前是实现细节的内容固化为接口的永久特性,例如 vectorize 是作为类而不是闭包实现的事实,或者方法是直接实现还是使用描述符实现。

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

讨论#

本提案的各种替代方案已在几个GitHub问题中讨论过

  1. pydata/sparse #1

  2. numpy/numpy #11129

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

本提案本身的详细讨论可在邮件列表和相关的拉取请求(1, 2, 3)中找到。