NEP 13 — 一种覆盖 Ufunc 的机制#

作者:

Blake Griffith

联系:

blake.g@utexas.edu

日期:

2013-07-10

作者:

Pauli Virtanen

作者:

Nathaniel Smith

作者:

Marten van Kerkwijk

作者:

Stephan Hoyer

日期:

2017-03-31

状态:

最终

更新:

2023-02-19

作者:

Roy Smart

执行摘要#

NumPy 的通用函数 (ufunc) 目前具有一些有限的功能,可用于使用 __array_prepare____array_wrap__ [1]ndarray 的用户定义子类进行操作,并且对任意对象几乎没有支持。例如 SciPy 的稀疏矩阵 [2] [3]

在这里,我们建议添加一种基于 ufunc 检查其每个参数是否存在 __array_ufunc__ 方法来覆盖 ufunc 的机制。如果发现 __array_ufunc__,则 ufunc 将操作交给该方法。

这涵盖了 Travis Oliphant 使用多方法对 NumPy 进行改造的提案 [4] 的一部分内容,该提案将解决相同的问题。此处的机制更紧密地遵循 Python 允许类覆盖 __mul__ 和其他二元运算符的方式。它还专门解决了二元运算符和 ufunc 如何交互的问题。(请注意,在早期版本中,覆盖称为 __numpy_ufunc__。已经实现了该功能,但行为并不完全正确,因此更改了名称。)

下面描述的 __array_ufunc__ 要求任何相应的 Python 二元运算符(__mul__ 等)都应以特定方式实现,并与 NumPy 的 ndarray 语义兼容。不满足此条件的对象无法覆盖任何 NumPy ufunc。我们没有指定将来兼容的路径来放宽此要求——此处任何更改都需要对第三方代码进行相应的更改。

动机#

人们普遍认为,当前用于分发 Ufunc 的机制不足。已经进行了长时间的讨论和其他提出的解决方案 [5][6]

使用 ndarray 的子类与 ufunc 的使用仅限于 __array_prepare____array_wrap__ 来准备输出参数,但这不允许您例如更改参数的形状或数据。尝试对不继承自 ndarray 的对象使用 ufunc 更加困难,因为输入参数往往会被转换为对象数组,这最终会产生令人惊讶的结果。

以 ufunc 与稀疏矩阵的互操作性为例。

In [1]: import numpy as np
import scipy.sparse as sp

a = np.random.randint(5, size=(3,3))
b = np.random.randint(5, size=(3,3))

asp = sp.csr_matrix(a)
bsp = sp.csr_matrix(b)

In [2]: a, b
Out[2]:(array([[0, 4, 4],
               [1, 3, 2],
               [1, 3, 1]]),
        array([[0, 1, 0],
               [0, 0, 1],
               [4, 0, 1]]))

In [3]: np.multiply(a, b) # The right answer
Out[3]: array([[0, 4, 0],
               [0, 0, 2],
               [4, 0, 1]])

In [4]: np.multiply(asp, bsp).todense() # calls __mul__ which does matrix multi
Out[4]: matrix([[16,  0,  8],
                [ 8,  1,  5],
                [ 4,  1,  4]], dtype=int64)

In [5]: np.multiply(a, bsp) # Returns NotImplemented to user, bad!
Out[5]: NotImplemented

返回 NotImplemented 给用户是不应该发生的。此外

In [6]: np.multiply(asp, b)
Out[6]: array([[ <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>,
                    <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>,
                    <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>],
                   [ <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>,
                    <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>,
                    <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>],
                   [ <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>,
                    <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>,
                    <3x3 sparse matrix of type '<class 'numpy.int64'>'
                with 8 stored elements in Compressed Sparse Row format>]], dtype=object)

这里,稀疏矩阵似乎被转换为对象数组标量,然后与 b 数组的所有元素相乘。但是,这种行为比有用的更令人困惑,并且 TypeError 会更好。

本提案 *不会* 解决 scipy.sparse 矩阵的问题,这些矩阵的乘法语义与 NumPy 数组不兼容。但是,其目标是能够编写其他具有严格与 ndarray 兼容语义的自定义数组类型。

建议的接口#

标准数组类 ndarray 获得了一个 __array_ufunc__ 方法,对象可以通过覆盖此方法来覆盖 Ufunc(如果它们是 ndarray 的子类)或定义自己的方法。方法签名是

def __array_ufunc__(self, ufunc, method, *inputs, **kwargs)

这里

  • ufunc 是调用的 ufunc 对象。

  • method 是一个字符串,指示ufunc的调用方式,可以是 "__call__"(表示直接调用),也可以是它的其中一个方法:"reduce""accumulate""reduceat""outer""at"

  • inputs 是一个元组,包含ufunc的输入参数。

  • kwargs 包含传递给函数的任何可选参数或关键字参数。这包括任何out参数,这些参数始终包含在一个元组中。

因此,参数被规范化:只有必需的输入参数(inputs)作为位置参数传递,所有其他参数都作为关键字参数字典(kwargs)传递。特别是,如果有输出参数(不是None),则它们作为元组在out关键字参数中传递(即使对于reduceaccumulatereduceat方法,在当前所有情况下只有一个输出是有意义的)。

函数调度过程如下:

  • 如果输入、输出或where参数之一实现了__array_ufunc__,则执行它而不是ufunc。

  • 如果多个参数实现了__array_ufunc__,则按以下顺序尝试:子类优先于超类,输入优先于输出,输出优先于where,否则从左到右。

  • 第一个__array_ufunc__方法返回非NotImplemented的值将决定ufunc的返回值。

  • 如果所有输入参数的__array_ufunc__方法都返回NotImplemented,则会引发TypeError

  • 如果__array_ufunc__方法引发错误,则立即传播该错误。

  • 如果没有任何输入参数具有__array_ufunc__方法,则执行将回退到默认的ufunc行为。

上述情况有一个前提:如果一个类具有__array_ufunc__属性,但它与ndarray.__array_ufunc__相同,则忽略该属性。这对于ndarray的实例以及没有覆盖其继承的__array_ufunc__实现的ndarray子类来说都是如此。

类型转换层次结构#

Python 运算符重载机制在如何编写重载方法方面提供了很大的自由度,为了获得可预测的结果,需要一定的规范。在这里,我们讨论一种理解一些含义的方法,这可以为设计提供输入。

明确哪些类型可以“向上转换”到其他类型(可能是间接的,例如,实现了间接 A->B->C 但没有直接 A->C)是有用的。如果__array_ufunc__的实现遵循一致的类型转换层次结构,则可以用来理解操作的结果。

类型转换可以用如下定义的来表示:

对于每个__array_ufunc__方法,从每个可能的输入类型到每个可能的输出类型绘制有向边。

也就是说,在y = x.__array_ufunc__(a, b, c, ...)返回非NotImplemented或引发错误的每种情况下,绘制边type(a) -> type(y)type(b) -> type(y)……

如果生成的图是无环的,它定义了一个一致的类型转换层次结构(类型之间明确的部分排序)。在这种情况下,涉及多种类型的操作通常会可预测地产生“最高”类型的结果,或者引发TypeError。参见本节末尾的示例。

如果图有环,则__array_ufunc__类型转换未定义,并且诸如type(multiply(a, b)) != type(multiply(b, a))type(add(a, add(b, c))) != type(add(add(a, b), c))的情况并不会被排除(然后可能总是可能)。

如果类型转换层次结构定义良好,对于每个类 A,所有其他定义了__array_ufunc__的类都属于三组中的一组:

  • 高于 A:A 可以(间接)向上转换为的类型。

  • 低于 A:可以(间接)向上转换为 A 的类型。

  • 不兼容:既不高于也不低于 A;任何(间接)向上转换都不可能的类型。

请注意,NumPy ufunc 的遗留行为是尝试通过np.asarray()将未知对象转换为ndarray。这相当于在图中将ndarray放在这些对象之上。由于我们在上面定义了ndarray对于具有自定义__array_ufunc__的类返回NotImplemented,这使得ndarray在类型层次结构中低于这些类,从而允许重载操作。

鉴于上述情况,描述传递操作的二元ufunc应该旨在定义一个定义良好的转换层次结构。这对于所有ufunc来说可能也是一种合理的方法——对此的例外情况应该仔细考虑是否会产生任何意外的行为。

示例

类型转换层次结构。

_images/nep0013_image1.png

类型 A 的__array_ufunc__可以处理返回 C 的ndarray,B 可以处理返回 B 的ndarray 和 D,C 可以处理返回 C 的 A 和 B,但不能处理 ndarray 或 D。结果是一个有向无环图,并定义了一个类型转换层次结构,关系为C > AC > ndarrayC > B > ndarrayC > B > D。类型 A 与 B、D、ndarray 不兼容,D 与 A 和 ndarray 不兼容。涉及这些类的ufunc表达式应该产生所涉及的最高类型的结果,或者引发TypeError

示例

__array_ufunc__图中的一个环。

_images/nep0013_image2.png

在这种情况下,__array_ufunc__关系具有长度为 1 的环,并且不存在类型转换层次结构。二元运算不是可交换的:type(a + b) is A,但type(b + a) is B

示例

__array_ufunc__图中的较长环。

_images/nep0013_image3.png

在这种情况下,__array_ufunc__ 关系具有更长的循环,并且不存在类型转换层次结构。二元运算仍然是可交换的,但是类型传递性丢失了:type(a + (b + c)) is A,但是 type((a + b) + c) is C

子类层次结构#

通常,希望在ufunc类型转换层次结构中镜像类层次结构。建议是,除非输入是同一类或超类的实例,否则类的__array_ufunc__ 实现通常应返回NotImplemented。这保证了在类型转换层次结构中,超类在下方,子类在上,其他类是不兼容的。对此的例外需要检查它们是否符合隐式类型转换层次结构。

注意

请注意,此处定义的类型转换层次结构和类层次结构的方向“相反”。原则上,让__array_ufunc__ 也处理子类的实例也是一致的。在这种情况下,“子类优先”调度规则将确保相对相似的结果。但是,行为的规定就不那么明确了。

如果方法一致地使用super() 传递类层次结构[7],则可以轻松构建子类。为此,ndarray 有其自己的__array_ufunc__ 方法,相当于

def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
    # Cannot handle items that have __array_ufunc__ (other than our own).
    outputs = kwargs.get('out', ())
    objs = inputs + outputs
    if "where" in kwargs:
        objs = objs + (kwargs["where"], )
    for item in objs:
        if (hasattr(item, '__array_ufunc__') and
                type(item).__array_ufunc__ is not ndarray.__array_ufunc__):
            return NotImplemented

    # If we didn't have to support legacy behaviour (__array_prepare__,
    # __array_wrap__, etc.), we might here convert python floats,
    # lists, etc, to arrays with
    # items = [np.asarray(item) for item in inputs]
    # and then start the right iterator for the given method.
    # However, we do have to support legacy, so call back into the ufunc.
    # Its arguments are now guaranteed not to have __array_ufunc__
    # overrides, and it will do the coercion to array for us.
    return getattr(ufunc, method)(*items, **kwargs)

请注意,作为特例,即使对于具有未覆盖默认ndarray 实现的ndarray 子类,ufunc 分派机制也不会调用此ndarray.__array_ufunc__ 方法。因此,调用ndarray.__array_ufunc__ 不会导致嵌套的 ufunc 分派循环。

super() 的使用对于仅添加属性(如单位)的ndarray 子类特别有用。在其__array_ufunc__ 实现中,此类类可以对其自身相关的参数进行可能的调整,并使用super() 传递到超类实现,直到实际完成 ufunc,然后对输出进行可能的调整。

通常,__array_ufunc__ 的自定义实现应避免嵌套分派循环,在这种循环中,一个不仅通过getattr(ufunc, method)(*items, **kwargs) 调用 ufunc,而且还捕获可能的异常等。与往常一样,可能会有例外。例如,对于像MaskedArray 这样的类,它只关心它包含的内容是ndarray 子类,使用__array_ufunc__ 的重新实现可能更容易通过直接将 ufunc 应用于其数据,然后调整掩码来完成。实际上,可以认为这是类的一部分,它决定它是否可以处理另一个参数(即它在类型层次结构中的位置)。在这种情况下,如果试验失败,则应返回NotImplemented。因此,实现将类似于

def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
    # for simplicity, outputs are ignored here.
    unmasked_items = tuple((item.data if isinstance(item, MaskedArray)
                            else item) for item in inputs)
    try:
        unmasked_result = getattr(ufunc, method)(*unmasked_items, **kwargs)
    except TypeError:
        return NotImplemented
    # for simplicity, ignore that unmasked_result could be a tuple
    # or a scalar.
    if not isinstance(unmasked_result, np.ndarray):
        return NotImplemented
    # now combine masks and view as MaskedArray instance
    ...

作为一个具体的例子,考虑一个数量和一个掩码数组类,它们都覆盖了__array_ufunc__,具有特定实例qma,后者包含一个常规数组。执行np.multiply(q, ma),ufunc 将首先分派到q.__array_ufunc__,后者返回NotImplemented(因为数量类将自身转换为数组并调用super(),它传递到ndarray.__array_ufunc__,后者看到ma 上的覆盖)。接下来,ma.__array_ufunc__ 获得机会。它不知道数量,如果它也返回NotImplemented,则会产生TypeError。但在我们的示例实现中,它使用getattr(ufunc, method) 来有效地计算np.multiply(q, ma.data)。这将再次传递到q.__array_ufunc__,但是这一次,由于ma.data 是一个常规数组,它将返回一个也是数量的结果。由于这是ndarray 的子类,ma.__array_ufunc__ 可以将其转换为掩码数组,从而返回结果(显然,如果它不是数组子类,它仍然可以返回NotImplemented)。

请注意,在上文中讨论的类型层次结构中,这是一个有点棘手的例子,因为MaskedArray 的位置很奇怪:它位于ndarray 的所有子类之上,因为它可以将它们转换为它自己的类型,但它本身不知道如何在 ufunc 中与它们交互。

关闭 Ufunc#

对于某些类,Ufunc 没有意义,并且像某些其他特殊方法(例如__hash____iter__[8] 一样,可以通过将__array_ufunc__ 设置为None 来指示 Ufunc 不可用。如果在设置__array_ufunc__ = None 的任何操作数上调用 Ufunc,它将无条件地引发TypeError

在类型转换层次结构中,这明确表示该类型相对于ndarray 是不兼容的。

与 Python 的二元运算符结合使用时的行为#

ndarray 中的 Python 运算符重载机制与__array_ufunc__ 机制相结合。对于 Python 用于实现二元运算(如*+)的特殊方法调用,例如ndarray.__mul__(self, other),NumPy 的ndarray 实现以下行为

  • 如果other.__array_ufunc__ is Nonendarray 返回NotImplemented。控制权返回给 Python,后者将依次尝试在other 上调用相应的反身方法(例如other.__rmul__)(如果存在)。

  • 如果other 上缺少__array_ufunc__ 属性,并且other.__array_priority__ > self.__array_priority__ndarray 也返回NotImplemented(逻辑与前一种情况一样)。这确保了与旧版 NumPy 的向后兼容性。

  • 否则,ndarray 将单方面调用相应的 Ufunc。Ufunc 永远不会返回NotImplemented,因此**诸如**other.__rmul__**之类的反身方法不能用于覆盖 NumPy 数组的算术运算,如果**__array_ufunc__** 设置为**None**以外的任何值。**相反,需要通过以与相应 Ufunc(例如np.multiply)一致的方式实现__array_ufunc__ 来更改它们的行为。有关受影响的运算符及其对应的 ufunc 列表,请参见运算符和 NumPy Ufunc 列表

如果某个类希望修改与ndarray在二元运算中的交互方式,则有两个选择:

  1. 实现__array_ufunc__并遵循NumPy的Python二元运算语义(见下文)。

  2. 设置__array_ufunc__ = None,并自由实现Python二元运算。在这种情况下,在此参数上调用的ufunc将引发TypeError(参见关闭Ufuncs)。

实现二元运算的建议#

对于大多数数值类,覆盖二元运算最简单的方法是定义__array_ufunc__并覆盖相应的Ufunc。然后,该类可以像ndarray本身一样,根据Ufuncs定义二元运算符。在这里,必须注意确保允许其他类指示它们不兼容,即实现应该类似于:

def _disables_array_ufunc(obj):
    try:
        return obj.__array_ufunc__ is None
    except AttributeError:
        return False

class ArrayLike:
    ...
    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        ...
        return result

    # Option 1: call ufunc directly
    def __mul__(self, other):
        if _disables_array_ufunc(other):
            return NotImplemented
        return np.multiply(self, other)

    def __rmul__(self, other):
        if _disables_array_ufunc(other):
            return NotImplemented
        return np.multiply(other, self)

    def __imul__(self, other):
        return np.multiply(self, other, out=(self,))

    # Option 2: call into one's own __array_ufunc__
    def __mul__(self, other):
        return self.__array_ufunc__(np.multiply, '__call__', self, other)

    def __rmul__(self, other):
        return self.__array_ufunc__(np.multiply, '__call__', other, self)

    def __imul__(self, other):
        result = self.__array_ufunc__(np.multiply, '__call__', self, other,
                                      out=(self,))
        if result is NotImplemented:
            raise TypeError(...)

为了了解为什么需要谨慎,请考虑另一个类other,它不知道如何处理数组和ufunc,因此已将__array_ufunc__设置为None,但知道如何进行乘法:

class MyObject:
    __array_ufunc__ = None
    def __init__(self, value):
        self.value = value
    def __repr__(self):
        return "MyObject({!r})".format(self.value)
    def __mul__(self, other):
        return MyObject(1234)
    def __rmul__(self, other):
        return MyObject(4321)

对于上述任一选项,我们都得到预期结果。

mine = MyObject(0)
arr = ArrayLike([0])

mine * arr    # -> MyObject(1234)
mine *= arr   # -> MyObject(1234)
arr * mine    # -> MyObject(4321)
arr *= mine   # -> TypeError

在这里,在第一个和第二个示例中,调用了mine.__mul__(arr),结果立即得到。在第三个示例中,首先调用arr.__mul__(mine)。在选项(1)中,对mine.__array_ufunc__ is None的检查将成功,因此返回NotImplemented,这导致执行mine.__rmul__(arg)。在选项(2)中,它可能是在arr.__array_ufunc__内部,很明显无法处理另一个参数,再次返回NotImplemented,从而导致控制权传递给mine.__rmul__

对于第四个示例(就地运算符),我们在这里遵循了ndarray并确保我们从不返回NotImplemented,而是引发TypeError。在选项(1)中,这是间接发生的:我们传递给np.multiply,后者又立即引发TypeError,因为它的一个操作数(out[0])禁用了Ufuncs。在选项(2)中,我们直接传递给arr.__array_ufunc__,它将返回NotImplemented,我们捕获它。

注意

不允许就地运算返回NotImplemented的原因是这些不能通过简单的反向操作来普遍替换:大多数数组操作假设实例的内容已就地更改,并且不期望一个新实例。此外,ndarr[:] *= mine意味着什么?假设它意味着ndarr[:] = ndarr[:] * mine,就像python在ndarr.__imul__返回NotImplemented时默认情况下所做的那样,这很可能是错误的。

现在考虑一下如果我们没有添加检查会发生什么。对于选项(1),相关的情况是我们没有检查__array_func__是否设置为None。在第三个示例中,调用了arr.__mul__(mine),如果没有检查,这将转到np.multiply(arr, mine)。这尝试arr.__array_ufunc__,它返回NotImplemented并看到mine.__array_ufunc__ is None,因此引发TypeError

对于选项(2),相关的示例是第四个,其中arr *= mine:如果我们让NotImplemented通过,python将用arr = mine.__rmul__(arr)替换它,这是不希望的。

因为Ufunc覆盖和Python的二元运算的语义几乎相同,所以在大多数情况下,选项(1)和(2)将使用相同的__array_ufunc__实现产生相同的结果。一个例外是当第二个参数是第一个参数的子类时尝试实现的顺序,这是由于一个预计将在Python 3.7中修复的Python错误[9]

一般来说,我们建议采用选项(1),这是与ndarray本身使用的方式最相似的选项。请注意,选项(1)是病毒式的,这意味着任何其他希望支持与您的类进行二元运算的类现在也必须遵循这些规则才能支持与ndarray的二元算术运算(即,它们必须实现__array_ufunc__或将其设置为None)。我们认为这是一件好事,因为它确保了所有支持ufunc和算术的对象的一致性。

为了使实现这样的类数组类更容易,mixin类NDArrayOperatorsMixin为所有具有相应Ufuncs的二元运算符提供选项(1)样式的覆盖。希望为兼容版本的NumPy实现__array_ufunc__但也需要支持与旧版本NumPy数组进行二元算术运算的类应该确保__array_ufunc__也可以用来实现它们支持的所有二元运算。

最后,我们注意到我们广泛讨论了是否更有意义地要求像MyObject这样的类实现完整的__array_ufunc__ [6]。最终,允许类选择退出是更可取的,上述推理使我们同意对ndarray本身进行类似的实现。选择退出机制需要禁用Ufuncs,因此类不能定义Ufuncs以返回与相应的二元运算不同的结果(即,如果定义了np.add(x, y),它应该与x + y匹配)。我们的目标是通过使Python的调度规则或NumPy的调度规则能够使用,而不是同时使用两者的混合,来尽可能简化与NumPy数组的二元运算的调度逻辑。

运算符和NumPy Ufuncs列表#

这是一个完整的Python二元运算符列表以及ndarrayNDArrayOperatorsMixin使用的相应NumPy Ufuncs。

符号

运算符

NumPy Ufunc(s)

<

lt

less()

<=

le

less_equal()

==

eq

equal()

!=

ne

not_equal()

>

gt

greater()

>=

ge

greater_equal()

+

add

add()

-

sub

subtract()

*

mul

multiply()

/

truediv (Python 3)

true_divide()

/

div (Python 2)

除法 (divide)

//

向下取整除法 (floordiv)

向下取整除法 (floor_divide)

%

取模 (mod)

取余 (remainder)

不适用 (NA)

商和余数 (divmod)

商和余数 (divmod)

**

幂运算 (pow)

power() [10]

<<

左移 (lshift)

左移 (left_shift)

>>

右移 (rshift)

右移 (right_shift)

&

按位与 (and_)

按位与 (bitwise_and)

^

按位异或 (xor_)

按位异或 (bitwise_xor)

|

按位或 (or_)

按位或 (bitwise_or)

@

矩阵乘法 (matmul)

尚未实现为通用函数 (ufunc) [11]

以下是单目运算符列表

符号

运算符

NumPy Ufunc(s)

-

取负 (neg)

取负 (negative)

+

取正 (pos)

positive() [12]

不适用 (NA)

取绝对值 (abs)

取绝对值 (absolute)

~

按位取反 (invert)

按位取反 (invert)

未来扩展到其他函数#

一些 NumPy 函数可以实现为(广义)ufunc,在这种情况下,可以通过 __array_ufunc__ 方法覆盖它们。一个主要的候选者是 matmul(),它目前不是 ufunc,但可以相对容易地重写为(一组)广义 ufunc。类似的情况也可能发生在诸如 median()min()argsort() 等函数上。