NEP 43 — 增强 UFuncs 的扩展性#

标题:

增强 UFuncs 的扩展性

作者:

Sebastian Berg

状态:

草稿

类型:

标准

创建:

2020-06-20

注意

此 NEP 是一个系列中的第四个

  • NEP 40 解释了 NumPy 的 dtype 实现的缺点。

  • NEP 41 概述了我们提出的替代方案。

  • NEP 42 描述了新设计的与数据类型相关的 API。

  • NEP 43(本文档)描述了新设计的通用函数 API。

摘要#

之前的 NEP 42 提出创建新的 DType,这些 DType 可以由 NumPy 本身之外的用户定义。NEP 42 的实现将使用户能够创建具有自定义 dtype 和存储值的数组。本 NEP 概述了 NumPy 将如何处理将来具有自定义 dtype 的数组。在 NumPy 数组上操作的最重要的函数是所谓的“通用函数”(ufunc),其中包括所有数学函数,例如 np.addnp.multiply,甚至 np.matmul。这些 ufunc 必须能够高效地处理具有不同数据类型的多个数组。

本 NEP 提出扩展 ufunc 的设计。它在可以处理许多不同 dtype(如浮点数或整数)的 ufunc 和定义特定 dtype 的高效操作的新 ArrayMethod 之间做出了新的区分。

注意

私有和外部 API 的详细信息可能会更改以反映用户评论和实现约束。底层原理和选择不应发生重大变化。

动机和范围#

本 NEP 的目标是扩展通用函数以支持 NEP 41 和 42 中详细介绍的新 DType 系统。虽然主要动机是启用新的用户定义的 DType,但这也将大大简化为 NumPy 字符串或结构化 DType 定义通用函数。到目前为止,由于其参数特性(比较 NEP 41 和 42),例如字符串长度,这些 DType 不受任何 NumPy 函数(如 np.addnp.equal)的支持。

数组上的函数必须处理许多不同的步骤,这些步骤在“UFunc 调用中涉及的步骤”部分中有更详细的描述。最重要的步骤是

  • 组织定义特定 DType 的 ufunc 调用的所有功能。这通常称为“内循环”。

  • 处理找不到完全匹配实现的输入。例如,当添加 int32float64 时,int32 会转换为 float64。这需要一个不同的“提升”步骤。

在组织和定义这些步骤之后,我们需要

  • 定义用于自定义上述两点的用户 API。

  • 允许方便地重用现有功能。例如,表示物理单位(例如米)的 DType 应该能够回退到 NumPy 现有的数学实现。

本 NEP 详细说明了如何在 NumPy 中实现这些要求

  • 目前 ufunc 定义中所有特定于 DTyper 的功能都将定义为新的 ArrayMethod 对象的一部分。此 ArrayMethod 对象将是描述在数组上操作的任何函数的新首选方法。

  • Ufuncs 将分派到 ArrayMethod 并可能使用提升来查找要使用的正确 ArrayMethod。这将在 提升和分派 部分中进行描述。

每个部分都将概述一个新的 C-API。预计未来的 Python API 将非常相似,并且为了易读性,C-API 将以 Python 代码的形式呈现。

NEP 提出对 NumPy ufunc 内部进行大量但必要的重构。此现代化不会直接影响最终用户,并且不仅是新 DType 的必要步骤,而且本身也是一项维护工作,预计将有助于未来对 ufunc 机器的改进。

虽然提出的最重要的重构是新的 ArrayMethod 对象,但最大的长期考虑因素是提升和分派的 API 选择。

向后兼容性#

NEP 41 中也列出了普遍的向后兼容性问题。

绝大多数用户除了 NumPy 版本的典型更改之外不应看到任何更改。有三个主要用户或用例受到建议更改的影响

  1. Numba 包直接访问 NumPy C 循环并直接修改 NumPy ufunc 结构以用于其自身目的。

  2. Astropy 使用它自己的“类型解析器”,这意味着从现有的类型解析到新的默认提升器的默认切换需要谨慎。

  3. 目前可以为 dtype 实例注册循环。这在理论上对结构化 dtype 有用,并且是此处提出的 DType 解析步骤之后发生的解析步骤。

此 NEP 将尽力保持尽可能多的向后兼容性。但是,这两个项目都表示愿意适应重大更改。

NumPy 能够提供向后兼容性的主要原因是

  • 可以包装现有的内循环,在调用中添加间接层,但保持完全向后兼容性。在这种情况下,get_loop 函数可以搜索现有的内循环函数(直接存储在 ufunc 上),以便即使在潜在的直接结构访问的情况下也能保持完全兼容性。

  • 可以调用旧版类型解析器作为回退(可能缓存结果)。可能需要调用两次解析器(一次用于 DType 解析,一次用于 resolve_descriptor 实现)。

  • 在大多数情况下,回退到旧版类型解析器应该处理为这种结构化 dtype 实例定义的循环。这是因为如果没有其他 np.Void 实现,旧版回退至少在最初会保留旧行为。

掩码类型解析器特别不会保留支持,但没有已知的用户(包括 NumPy 本身,它只使用默认版本)。

此外,不会尝试与调用保持兼容,而不是提供正常或掩码类型解析器。因为 NumPy 只会将其用作回退。没有已知用户使用这种(未记录的)可能性。

虽然上述更改可能会破坏某些工作流程,但我们认为长期改进远远超过了这一点。此外,像 astropy 和 Numba 这样的包能够适应,因此最终用户可能需要更新他们的库,但不需要更新他们的代码。

使用和影响#

此 NEP 重构了 NumPy 数组操作在 NumPy 内部和外部实现者中定义的方式。NEP 主要涉及那些为自定义 DType 扩展 ufunc 或创建自定义 ufunc 的人。它并非旨在完成所有潜在的用例,而是重构 NumPy 以使其可扩展,并允许在出现新的问题或功能请求时解决它们。

概述和最终用户 API#

为了概述本 NEP 如何提出结构化 ufunc,以下描述了提议的重构对最终用户的潜在影响。

当考虑只有一个输入的 ufunc 时,通用函数很像在数组的 DType 上定义的 Python 方法

res = np.positive(arr)

可以实现(概念上)为

positive_impl = arr.dtype.positive
res = positive_impl(arr)

但是,与方法不同,positive_impl 并不存储在 dtype 本身中。它实际上是 np.positive 对特定 DType 的实现。当前的 NumPy 部分地公开了此“实现选择”,使用通用函数中的 dtype(或更准确的 signature)属性,尽管这些属性很少使用

np.positive(arr, dtype=np.float64)

强制 NumPy 使用专门为 Float64 DType 编写的 positive_impl

此 NEP 通过创建一个新对象来表示 positive_impl,使这种区别更加明确

positive_impl = np.positive.resolve_impl((type(arr.dtype), None))
# The `None` represents the output DType which is automatically chosen.

虽然 positive_impl 对象和 resolve_impl 方法的创建是本 NEP 的一部分,但以下代码

res = positive_impl(arr)

可能不会最初实现,并且不是重新设计的核心。

通常,NumPy 通用函数可以接收许多输入。这需要通过考虑所有输入来查找实现,并使 ufunc 成为相对于输入 DType 的“多方法”

add_impl = np.add.resolve_impl((type(arr1.dtype), type(arr2.dtype), None))

此 NEP 定义了 positive_impladd_impl 如何表示为新的 ArrayMethod,该方法可以在 NumPy 之外实现。此外,它定义了 resolve_impl 如何实现和解决分派和提升。

在查看 UFunc 调用中涉及的步骤 部分后,这些拆分的原因可能会更加清楚。

定义新的 ufunc 实现#

以下是关于如何将新的实现(在本例中是定义字符串相等性)添加到 ufunc 的模拟。

class StringEquality(BoundArrayMethod):
    nin = 1
    nout = 1
    # DTypes are stored on the BoundArrayMethod and not on the internal
    # ArrayMethod, to reference cyles.
    DTypes = (String, String, Bool)

    def resolve_descriptors(self: ArrayMethod, DTypes, given_descrs):
        """The strided loop supports all input string dtype instances
        and always returns a boolean. (String is always native byte order.)

        Defining this function is not necessary, since NumPy can provide
        it by default.

        The `self` argument here refers to the unbound array method, so
        that DTypes are passed in explicitly.
        """
        assert isinstance(given_descrs[0], DTypes[0])
        assert isinstance(given_descrs[1], DTypes[1])
        assert given_descrs[2] is None or isinstance(given_descrs[2], DTypes[2])

        out_descr = given_descrs[2]  # preserve input (e.g. metadata)
        if given_descrs[2] is None:
            out_descr = DTypes[2]()

        # The operation is always "no" casting (most ufuncs are)
        return (given_descrs[0], given_descrs[1], out_descr), "no"

    def strided_loop(context, dimensions, data, strides, innerloop_data):
        """The 1-D strided loop, similar to those used in current ufuncs"""
        # dimensions: Number of loop items and core dimensions
        # data: Pointers to the array data.
        # strides: strides to iterate all elements
        n = dimensions[0]  # number of items to loop over
        num_chars1 = context.descriptors[0].itemsize
        num_chars2 = context.descriptors[1].itemsize

        # C code using the above information to compare the strings in
        # both arrays.  In particular, this loop requires the `num_chars1`
        # and `num_chars2`.  Information which is currently not easily
        # available.

np.equal.register_impl(StringEquality)
del StringEquality  # may be deleted.

此定义足以创建一个新的循环,并且该结构允许将来进行扩展;这是 NumPy 本身实现强制转换所需的。我们在此处使用 BoundArrayMethodcontext 结构。这些将在后面详细描述和说明其动机。简而言之

  • context 是 Python 传递给其方法的 self 的泛化。

  • BoundArrayMethod 等价于 Python 中的区分,即 class.method 是一个方法,而 class().method 返回一个“绑定”方法。

自定义分派和提升#

当调用 np.positive.resolve_impl() 时,查找正确的实现很大程度上是一个实现细节。但是,在某些情况下,当没有实现完全匹配请求的 DType 时,可能需要影响此过程。

np.multiple.resolve_impl((Timedelta64, Int8, None))

将没有完全匹配项,因为 NumPy 仅具有将 Timedelta64Int64 相乘的实现。在简单的情况下,NumPy 将使用默认的提升步骤尝试查找正确的实现,但要实现上述步骤,我们将允许以下操作

def promote_timedelta_integer(ufunc, dtypes):
    new_dtypes = (Timdelta64, Int64, dtypes[-1])
    # Resolve again, using Int64:
    return ufunc.resolve_impl(new_dtypes)

np.multiple.register_promoter(
    (Timedelta64, Integer, None), promote_timedelta_integer)

其中 Integer 是一个抽象 DType(参见 NEP 42)。

UFunc 调用涉及的步骤#

在深入探讨更详细的 API 选择之前,回顾 NumPy 中通用函数调用的步骤会有所帮助。

UFunc 调用分为以下步骤

  1. 处理 __array_ufunc__ 协议

  2. 提升和分派

    • 给定所有输入的 DType,找到正确的实现。例如,针对 float64int64 或用户定义的 DType 的实现。

    • 当不存在精确的实现时,必须执行提升。例如,将 float32float64 相加是通过首先将 float32 转换为 float64 来实现的。

  3. 参数化 dtype 解析

    • 一般来说,只要输出 DType 是参数化的,就必须找到(解析)参数。

    • 例如,如果循环添加两个字符串,则必须定义正确的输出(以及可能输入)dtype。 S5 + S4 -> S9,而 upper 函数的签名为 S5 -> S5

    • 当它们不是参数化时,将提供一个默认实现,该实现填充默认的 dtype 实例(例如,确保本地字节顺序)。

  4. 准备迭代

    • 此步骤主要由内部的 NpyIter(迭代器)处理。

    • 分配执行转换所需的所有输出和临时缓冲区。这需要在步骤 3 中解析的 dtype。

    • 找到最佳迭代顺序,其中包括有效实现广播的信息。例如,将单个值添加到数组会重复相同的值。

  5. 设置和获取 C 级函数

    • 如有必要,分配临时工作空间。

    • 查找 C 实现的轻量级内部循环函数。查找内部循环函数可以在未来允许专门的实现。例如,转换当前优化连续转换,并且约简具有当前在内部循环函数本身中处理的优化。

    • 发出内部循环是否需要 Python API 或是否可以释放 GIL(以允许线程)的信号。

    • 清除浮点异常标志。

  6. 执行实际计算

    • 运行特定于 DType 的内部循环函数。

    • 内部循环可能需要访问其他数据,例如在上一步骤中设置的 dtype 或其他数据。

    • 内部循环函数可能会被调用不定次数。

  7. 完成

    • 释放步骤 5 中分配的任何临时工作空间。

    • 检查浮点异常标志。

    • 返回结果。

ArrayMethod 提供了一个概念来对步骤 3 到 6 以及部分 7 进行分组。但是,新 ufunc 或 ArrayMethod 的实现者通常不需要自定义步骤 4 或 6 中的行为,NumPy 可以并且确实提供了这些行为。对于 ArrayMethod 实现者来说,需要自定义的核心步骤是步骤 3 和步骤 5。这些提供了自定义的内部循环函数以及潜在的特定于内部循环的设置。进一步的自定义是可能的,并且预计作为未来的扩展。

步骤 2 是提升和分派,并将使用新的 API 重构,以便在必要时允许自定义流程。

步骤 1 出于完整性而列出,不受此 NEP 的影响。

以下草图概述了步骤 2 到 6,重点介绍了如何处理 dtype 以及哪些部分是可自定义的(“已注册”)以及哪些部分由 NumPy 处理

_images/nep43-sketch.svg

ArrayMethod#

此 NEP 的核心提议是创建 ArrayMethod 作为描述特定于给定 DType 集的每个实现的对象。我们使用 class 语法来描述创建新的 ArrayMethod 对象所需的信息

class ArrayMethod:
    name: str  # Name, mainly useful for debugging

    # Casting safety information (almost always "safe", necessary to
    # unify casting and universal functions)
    casting: Casting = "no"

    # More general flags:
    flags: int

    def resolve_descriptors(self,
            Tuple[DTypeMeta], Tuple[DType|None]: given_descrs) -> Casting, Tuple[DType]:
        """Returns the safety of the operation (casting safety) and the
        """
        # A default implementation can be provided for non-parametric
        # output dtypes.
        raise NotImplementedError

    @staticmethod
    def get_loop(Context : context, strides, ...) -> strided_loop_function, flags:
        """Returns the low-level C (strided inner-loop) function which
        performs the actual operation.

        This method may initially private, users will be able to provide
        a set of optimized inner-loop functions instead:

        * `strided_inner_loop`
        * `contiguous_inner_loop`
        * `unaligned_strided_loop`
        * ...
        """
        raise NotImplementedError

    @staticmethod
    def strided_inner_loop(
            Context : context, data, dimensions, strides, innerloop_data):
        """The inner-loop (equivalent to the current ufunc loop)
        which is returned by the default `get_loop()` implementation."""
        raise NotImplementedError

使用 Context 提供有关函数调用的大部分静态信息

class Context:
    # The ArrayMethod object itself:
    ArrayMethod : method

    # Information about the caller, e.g. the ufunc, such as `np.add`:
    callable : caller = None
    # The number of input arguments:
    int : nin = 1
    # The number of output arguments:
    int : nout = 1
    # The actual dtypes instances the inner-loop operates on:
    Tuple[DType] : descriptors

    # Any additional information required. In the future, this will
    # generalize or duplicate things currently stored on the ufunc:
    #  - The ufunc signature of generalized ufuncs
    #  - The identity used for reductions

以及存储的属性 flags,用于指示

  • ArrayMethod 是否支持未对齐的输入和输出数组

  • 内部循环函数是否需要 Python API (GIL)

  • NumPy 是否必须检查浮点错误 CPU 标志。

注意:预计会根据需要添加更多信息。

调用 Context#

“上下文”对象类似于 Python 的 self,它被传递给所有方法。要理解为什么需要“上下文”对象及其内部结构,有助于记住 Python 方法可以用以下方式编写(另请参见 __get__ 的文档

class BoundMethod:
    def __init__(self, instance, method):
        self.instance = instance
        self.method = method

    def __call__(self, *args, **kwargs):
        return self.method.function(self.instance, *args, **kwargs)


class Method:
    def __init__(self, function):
        self.function = function

    def __get__(self, instance, owner=None):
        assert instance is not None  # unsupported here
        return BoundMethod(instance, self)

使用以下 method1method2,其行为相同

def function(self):
    print(self)

class MyClass:
    def method1(self):
        print(self)

    method2 = Method(function)

并且两者都将打印相同的结果

>>> myinstance = MyClass()
>>> myinstance.method1()
<__main__.MyClass object at 0x7eff65436d00>
>>> myinstance.method2()
<__main__.MyClass object at 0x7eff65436d00>

这里 self.instance 将是 Context 传递的所有信息。 Context 是一个泛化,并且必须传递其他信息

  • 与仅对单个类实例进行操作的方法不同,ArrayMethod 对多个输入数组以及多个 dtype 进行操作。

  • 上面 BoundMethod__call__ 只包含对函数的单个调用。但是,ArrayMethod 必须调用 resolve_descriptors,然后将该信息传递给内部循环函数。

  • Python 函数除了由其外部作用域定义的状态之外,没有其他状态。在 C 中,Context 能够根据需要提供其他状态。

就像 Python 需要区分方法和绑定方法一样,NumPy 将具有 BoundArrayMethod。这存储了构成 Context 部分的所有常量信息,例如

幸运的是,大多数用户甚至 ufunc 实现者都不必担心这些内部细节;就像很少有 Python 用户需要了解 __get__ dunder 方法一样。 Context 对象或 C 结构为快速 C 函数提供了所有必要的数据,并且 NumPy API 根据需要创建新的 ArrayMethodBoundArrayMethod

ArrayMethod 规范#

这些规范提供了最小的初始 C API,将来会扩展,例如允许专门的内部循环。

简而言之,NumPy 目前依赖于步幅内部循环,这将是最初定义 ufunc 的唯一允许方法。我们预计将来会添加 setup 函数或公开 get_loop

UFunc 需要与转换相同的信息,给出以下定义(另请参见 NEP 42 CastingImpl

  • 一个要传递给解析函数和内部循环的新结构

    typedef struct {
        PyObject *caller;  /* The ufunc object */
        PyArrayMethodObject *method;
    
        int nin, nout;
    
        PyArray_DTypeMeta **dtypes;
        /* Operand descriptors, filled in by resolve_desciptors */
        PyArray_Descr **descriptors;
    
        void *reserved;  // For Potential in threading (Interpreter state)
    } PyArrayMethod_Context
    

    此结构可能会附加到 NumPy 的未来版本中以包含其他信息,并且包含所有常量循环元数据。

    我们可以对该结构进行版本控制,尽管对 ArrayMethod 本身进行版本控制可能更简单。

  • 与转换类似,ufunc 可能需要查找正确的循环 dtype 或指示循环只能处理所涉及 DType 的某些实例(例如,仅本地字节顺序)。这由 resolve_descriptors 函数处理(与 CastingImplresolve_descriptors 相同)

    NPY_CASTING
    resolve_descriptors(
            PyArrayMethodObject *self,
            PyArray_DTypeMeta *dtypes,
            PyArray_Descr *given_dtypes[nin+nout],
            PyArray_Descr *loop_dtypes[nin+nout]);
    

    该函数根据给定的 given_dtypes 填充 loop_dtypes。这需要填写输出的描述符。通常还必须找到输入描述符,例如,在内部循环需要时确保本地字节顺序。

    在大多数情况下,ArrayMethod 将具有非参数化输出 DType,以便可以提供默认实现。

  • 附加的 void *user_data 通常会被类型化为扩展现有的 NpyAuxData * 结构

    struct {
        NpyAuxData_FreeFunc *free;
        NpyAuxData_CloneFunc *clone;
        /* To allow for a bit of expansion without breaking the ABI */
       void *reserved[2];
    } NpyAuxData;
    

    此结构目前主要用于 NumPy 内部转换机制,截至目前,必须提供 freeclone,尽管这可以放宽。

    与 NumPy 转换不同,绝大多数 ufunc 目前不需要此额外的暂存空间,但可能需要简单的标记功能,例如用于实现警告(请参见下面的错误和警告处理)。为了简化这一点,NumPy 将在未设置 user_data 时传递一个单一的零初始化的 npy_intp *注意,可以将其作为 Context 的一部分传递。

  • 可选的 get_loop 函数最初不会公开,以避免最终确定需要与转换一起进行设计选择的 API

    innerloop *
    get_loop(
        PyArrayMethod_Context *context,
        int aligned, int move_references,
        npy_intp *strides,
        PyArray_StridedUnaryOp **out_loop,
        NpyAuxData **innerloop_data,
        NPY_ARRAYMETHOD_FLAGS *flags);
    

    NPY_ARRAYMETHOD_FLAGS 可以指示是否需要 Python API 以及是否必须检查浮点错误。 move_references 当前在 NumPy 转换中用于内部使用。

  • 内部循环函数

    int inner_loop(PyArrayMethod_Context *context, ..., void *innerloop_data);
    

    将具有与当前内部循环相同的签名,并进行以下更改

    • 一个返回值,用于在返回 -1 而不是 0 时指示错误。在返回 -1 时,必须设置 Python 错误。

    • 新的第一个参数 PyArrayMethod_Context * 用于以方便的方式传入可能需要的有关 ufunc 或描述符的信息。

    • void *innerloop_data 将是 NpyAuxData **innerloop_data,由 get_loop 设置。如果 get_loop 未设置 innerloop_data,则会改为传递 npy_intp *(请参见下面的 错误处理 以了解原因)。

    注意:由于预计 get_loop 将是私有的,因此可以在最终公开之前修改 innerloop_data 的确切实现。

创建新的 BoundArrayMethod 将使用 PyArrayMethod_FromSpec() 函数。一个简写将允许使用 PyUFunc_AddImplementationFromSpec() 直接注册到 ufunc。规范预计将包含以下内容(这可能会在将来扩展)

typedef struct {
    const char *name;  /* Generic name, mainly for debugging */
    int nin, nout;
    NPY_CASTING casting;
    NPY_ARRAYMETHOD_FLAGS flags;
    PyArray_DTypeMeta **dtypes;
    PyType_Slot *slots;
} PyArrayMethod_Spec;

讨论和替代方案#

上面将 ArrayMethodContext 分开,以及对 BoundArrayMethod 的额外要求,是必要的拆分,反映了 Python 中方法和绑定方法的实现。

此要求的一个原因是它允许在许多情况下存储 ArrayMethod 对象,而无需保存对 DTypes 的引用,如果动态创建(和删除)DType,这可能很重要。(这是一个复杂的话题,在当前的 Python 中没有完整的解决方案,但该方法解决了转换方面的问题。)

这种结构似乎没有替代方案。将特定于 DType 的步骤与通用 ufunc 分发和提升分离对于允许未来的扩展和灵活性是绝对必要的。此外,它允许统一类型转换和 ufunc。

由于ArrayMethodBoundArrayMethod的结构将是不透明的并且可以扩展,因此除了选择将它们作为 Python 对象之外,几乎没有长期设计影响。

resolve_descriptors#

resolve_descriptors方法可能是此 NEP 的主要创新,它也是 NEP 42 中类型转换实现的核心。

通过确保每个ArrayMethod都提供resolve_descriptors,我们为UFunc 调用涉及的步骤中的步骤 3 定义了一个统一、清晰的 API。此步骤是分配输出数组所必需的,并且必须在准备类型转换之前发生。

虽然返回的类型转换安全性(NPY_CASTING)对于通用函数几乎总是“no”,但包含它有两个主要优点

  • -1表示发生了错误。如果设置了 Python 错误,则会引发该错误。如果未设置 Python 错误,则将其视为“不可能”的转换,并将设置自定义错误。(这种区别对于np.can_cast()函数很重要,该函数应该引发第一个错误并在第二种情况下返回False,对于典型的 ufunc 来说,它并不值得注意)。这一点正在考虑中,我们可能会使用 -1 表示一般错误,并使用不同的返回值表示不可能的转换。

  • 返回类型转换安全性对于 NEP 42 中的类型转换至关重要,并允许在其中未经修改地使用ArrayMethod

  • 将来可能需要实现快速但不安全的实现。例如,对于int64 + int64 -> int32,从类型转换的角度来看是不安全的。目前,这将使用int64 + int64 -> int64,然后转换为int32。跳过转换的实现必须表明它实际上包含“同类”转换,因此不被认为是“no”。

get_loop方法#

目前,NumPy ufunc 通常只提供一个步长循环,因此get_loop方法可能显得不必要。出于这个原因,我们计划最初将get_loop作为私有函数。

但是,get_loop是类型转换所必需的,即使在步长和连续循环之外,类型转换也使用专门的循环。因此,get_loop函数必须完全替代内部PyArray_GetDTypeTransferFunction

将来,get_loop可能会公开,或者公开一个新的setup函数以允许更多控制,例如允许分配工作内存。此外,我们可以扩展get_loop并允许ArrayMethod实现者不仅控制一维内部循环,还控制外部迭代。

扩展内部循环签名#

扩展内部循环签名是 NEP 的另一个核心且必要的组成部分。

传入 Context

传入Context可能允许通过向 context 结构添加新字段来扩展签名的未来。此外,它还提供对内部循环操作的数据类型实例的直接访问。对于参数化数据类型,这是必要的信息,例如比较两个字符串需要知道两个字符串的长度。 Context还可以保存可能很有用的信息,例如原始ufunc,这在报告错误时很有用。

原则上,传入 Context 是没有必要的,因为所有信息都可以包含在innerloop_data中并在get_loop函数中设置。在此 NEP 中,我们建议传递该结构以简化参数化 DType 循环的创建。

传入用户数据

当前的类型转换实现使用现有的NpyAuxData *来传入由get_loop定义的附加数据。当然,这种结构的使用还有其他替代方案,但它提供了一个简单的解决方案,该方案已在 NumPy 和公共 API 中使用。

NpyAyxData *是一个轻量级、已分配的结构,并且由于它已经存在于 NumPy 中用于此目的,因此似乎是一个自然的选择。为了简化某些用例(请参阅下面的“错误处理”),当未提供innerloop_data时,我们将改为传递npy_intp *innerloop_data = 0

注意:由于预计get_loop最初是私有的,因此我们可以在将其公开为公共 API 之前积累使用innerloop_data的经验。

返回值

指示错误的返回值是 NumPy 中一个重要的但目前缺失的功能。错误处理因 CPU 发出浮点错误的方式而进一步复杂化。两者将在下一节中讨论。

错误处理#

我们预计未来的内部循环通常会在发现错误后立即设置 Python 错误。当在没有锁定 GIL 的情况下运行内部循环时,这会变得复杂。在这种情况下,该函数将必须锁定 GIL,设置 Python 错误并返回-1以指示发生了错误:

int
inner_loop(PyArrayMethod_Context *context, ..., void *innerloop_data)
{
    NPY_ALLOW_C_API_DEF

    for (npy_intp i = 0; i < N; i++) {
        /* calculation */

        if (error_occurred) {
            NPY_ALLOW_C_API;
            PyErr_SetString(PyExc_ValueError,
                "Error occurred inside inner_loop.");
            NPY_DISABLE_C_API
            return -1;
        }
    }
    return 0;
}

浮点错误是特殊的,因为它们需要检查硬件状态,如果在内部循环函数本身内进行检查,则成本太高。因此,如果由ArrayMethod标记,NumPy 将处理这些错误。ArrayMethod如果标记不应检查这些错误,则绝不应该导致设置浮点错误标志。当调用多个函数时,这可能会发生干扰;特别是当类型转换是必要的时。

另一种解决方案是在以后的最终化步骤中只允许设置错误,届时 NumPy 也会检查浮点错误标志。

我们目前不建议使用此模式。它似乎更复杂,而且通常没有必要。虽然在循环中安全地获取 GIL 可能需要在将来传递额外的PyThreadStatePyInterpreterState(用于子解释器支持),但这可以接受并且可以预期。在稍后设置错误会增加复杂性:例如,如果操作被暂停(目前尤其可能发生在类型转换中),则每次发生这种情况时都需要显式运行错误检查。

我们预计立即设置错误是最简单、最方便的解决方案,更复杂的解决方案可能是未来的扩展。

处理警告稍微复杂一些:即使在简单情况下会多次发出警告,也应该为每个函数调用(即整个数组)准确地发出一次警告。为了简化此类用例,我们将默认传入npy_intp *innerloop_data = 0,可用于存储标志(或其他简单的持久数据)。例如,我们可以想象一个整数乘法循环,当发生溢出时发出警告

int
integer_multiply(PyArrayMethod_Context *context, ..., npy_intp *innerloop_data)
{
    int overflow;
    NPY_ALLOW_C_API_DEF

    for (npy_intp i = 0; i < N; i++) {
        *out = multiply_integers(*in1, *in2, &overflow);

        if (overflow && !*innerloop_data) {
            NPY_ALLOW_C_API;
            if (PyErr_Warn(PyExc_UserWarning,
                    "Integer overflow detected.") < 0) {
                NPY_DISABLE_C_API
                return -1;
            }
            *innerloop_data = 1;
            NPY_DISABLE_C_API
    }
    return 0;
}

待办事项:当未设置innerloop_data时传递npy_intp暂存空间的想法似乎很方便,但我对此不确定,因为我不知道任何类似的先例。“暂存空间”原则上也可以是context的一部分。

重用现有循环/实现#

对于许多 DType 来说,以上添加额外 C 级循环的定义将足够,并且只需要一个步长循环实现,如果循环适用于参数化 DType,则必须另外提供resolve_descriptors函数。

但是,在某些用例中,希望回调到现有实现。在 Python 中,这可以通过简单地调用原始 ufunc 来实现。

为了提高 C 的性能以及大型数组的性能,希望尽可能直接地重用现有的ArrayMethod,以便其内部循环函数可以直接使用,而无需额外的开销。因此,我们将允许从现有的ArrayMethod创建一个新的包装ArrayMethod

此包装的ArrayMethod将具有两个额外的方法

  • view_inputs(Tuple[DType]: input_descr) -> Tuple[DType]将用户输入描述符替换为与包装循环匹配的描述符。必须能够将输入视为输出。例如,对于Unit[Float64]("m") + Unit[Float32]("km"),这将返回float64 + int32。原始resolve_descriptors将将其转换为float64 + float64

  • wrap_outputs(Tuple[DType]: input_descr) -> Tuple[DType]将解析的描述符替换为所需的实际循环描述符。原始resolve_descriptors函数将在这两个调用之间被调用,因此输出描述符可能不会在第一个调用中设置。在上面的示例中,它将使用返回的float64(这可能会更改字节顺序),并进一步解析物理单位,从而形成最终签名

    Unit[Float64]("m") + Unit[Float64]("m") -> Unit[Float64]("m")
    

    UFunc 机制将负责将“km”输入转换为“m”。

view_inputs方法允许将正确的输入传递到原始resolve_descriptors函数中,而wrap_outputs确保使用正确的描述符进行输出分配和输入缓冲区转换。

这方面的一个重要用例是抽象 Unit DType,每个数值数据类型都有一个子类(可以动态创建)

Unit[Float64]("m")
# with Unit[Float64] being the concrete DType:
isinstance(Unit[Float64], Unit)  # is True

这样的Unit[Float64]("m")实例在类型提升方面具有明确定义的签名。 UnitDType 的作者可以通过包装现有的数学函数并使用上述两个额外的方法来实现大多数必要的逻辑。使用提升步骤,这将允许为抽象UnitDType 与ufunc创建一个注册一个提升器。然后,提升器可以在提升时动态添加包装的具体ArrayMethod,并且 NumPy 可以在第一次调用后缓存(或存储)它。

替代用例

另一个用例是Unit(float64, "m")DType,其中数值类型是 DType 参数的一部分。这种方法是可能的,但需要一个包装现有循环的自定义ArrayMethod。它还必须始终需要两个分发步骤(一个到UnitDType,另一个到数值类型)。

此外,有效的实现将需要能够从另一个ArrayMethod获取和重用内部循环函数。(这对于像 Numba 这样的用户来说可能是必要的,但尚不确定它是否应该是常见的模式,并且它本身无法从 Python 访问。)

提升和分发#

NumPy ufunc 从某种意义上说是多方法,因为它们同时操作(或使用)多个 DType。虽然输入(和输出)数据类型附加到 NumPy 数组,但ndarray类型本身不包含要应用于数据的功能的信息。

例如,给定输入

int_arr = np.array([1, 2, 3], dtype=np.int64)
float_arr = np.array([1, 2, 3], dtype=np.float64)
np.add(int_arr, float_arr)

必须找到执行操作的正确ArrayMethod。理想情况下,存在一个已定义的完全匹配,例如,对于np.add(int_arr, int_arr)ArrayMethod[Int64, Int64, out=Int64]完全匹配,可以使用。但是,对于np.add(int_arr, float_arr),没有直接匹配,需要提升步骤。

提升和分发过程#

通常,通过搜索所有输入 DType 的完全匹配来找到ArrayMethod。输出数据类型不应影响计算,但如果多个注册的ArrayMethod完全匹配,则将使用输出 DType 来找到更好的匹配。这将允许当前对于np.equal循环的区别,该循环同时定义了Object, Object -> Bool(默认)和Object, Object -> Object

最初,ArrayMethod 将仅为具体的 DType 定义,并且由于这些 DType 无法被子类化,因此可以保证精确匹配。将来,我们预计 ArrayMethod 也可为抽象的 DType 定义。在这种情况下,将找到最佳匹配,如下所述。

提升

如果存在匹配的 ArrayMethod,则分派将非常简单。但是,当不存在匹配的 ArrayMethod 时,需要额外的定义来实现此“提升”。

  • 默认情况下,任何 UFunc 都有一个提升,该提升使用所有输入的公共 DType 并再次分派。这对于大多数数学函数来说是明确定义的,但如果需要,可以禁用或自定义。例如,int32 + float64 会再次尝试使用 float64 + float64,这是公共的 DType。

  • 用户可以注册新的提升器,就像他们可以注册新的 ArrayMethod 一样。这些将使用抽象 DType 来允许匹配各种签名。提升函数的返回值应为新的 ArrayMethodNotImplemented。它必须在使用相同输入的多次调用中保持一致,以允许缓存结果。

提升函数的签名将是

promoter(np.ufunc: ufunc, Tuple[DTypeMeta]: DTypes): -> Union[ArrayMethod, NotImplemented]

请注意,DType 可能包含输出的 DType,但是,通常输出 DType 不会影响选择哪个 ArrayMethod

在大多数情况下,无需添加自定义提升函数。需要此函数的一个示例是与单位相乘:在 NumPy 中,timedelta64 可以乘以大多数整数,但 NumPy 仅为 timedelta64 * int64 定义了一个循环(ArrayMethod),因此与 int32 相乘将失败。

为了允许这样做,可以为 (Timedelta64, Integral, None) 注册以下提升器

def promote(ufunc, DTypes):
    res = list(DTypes)
    try:
        res[1] = np.common_dtype(DTypes[1], Int64)
    except TypeError:
        return NotImplemented

    # Could check that res[1] is actually Int64
    return ufunc.resolve_impl(tuple(res))

在这种情况下,就像需要 Timedelta64 * int64int64 * timedelta64 ArrayMethod 一样,将必须注册第二个提升器来处理整数首先传递的情况。

ArrayMethod 和提升器的分派规则

通过找到 DType 类层次结构定义的最佳匹配来发现提升器和 ArrayMethod。如果满足以下条件,则定义了最佳匹配

  • 签名与所有输入 DType 匹配,以便 issubclass(input_DType, registered_DType) 返回 true。

  • 没有其他提升器或 ArrayMethod 在任何输入方面更精确:issubclass(other_DType, this_DType) 为 true(这可能包括两者都相同的情况)。

  • 此提升器或 ArrayMethod 在至少一个输入或输出 DType 方面更精确。

如果返回 NotImplemented 或两个提升器对输入的匹配程度相同,则将发生错误。当现有提升器对于新功能不够精确时,必须添加新的提升器。为了确保此提升器优先,可能需要定义新的抽象 DType 作为现有 DType 的更精确的子类。

上述规则允许在提供输出或指定完整循环时进行专门化。这通常是不必要的,但允许解决 np.logic_or 等具有 Object, Object -> BoolObject, Object -> Object 循环的函数(默认使用第一个)。

讨论和替代方案#

除了解析并返回新的实现之外,我们还可以返回一组新的 DType 用于分派。但是,这也有一个缺点,即无法分派到在不同 ufunc 上定义的循环,也无法动态创建新的 ArrayMethod

被拒绝的替代方案

在上述内容中,提升器使用多重分派样式类型解析,而当前的 UFunc 机制使用第一个“安全”循环(另请参见 NEP 40)在一个有序的层次结构中。

虽然“安全”转换规则不够严格,但我们可以想象使用新的“提升”转换规则,或使用 common-DType 逻辑通过根据需要向上转换输入来找到最佳匹配循环。

这种方法的一个缺点是,仅向上转换允许将结果向上转换超出用户期望的范围:目前(将作为后备继续支持)任何仅定义 float64 循环的 ufunc 也将通过向上转换适用于 float16 和 float32

>>> from scipy.special import erf
>>> erf(np.array([4.], dtype=np.float16))  # float16
array([1.], dtype=float32)

并获得 float32 结果。无法更改 erf 函数以返回 float16 结果,而不更改以下代码的结果。总的来说,我们认为在可以定义精度较低的循环的情况下,不应发生自动向上转换,除非 ufunc 作者使用提升有意这样做。

此考虑意味着必须通过某种其他方法限制向上转换。

替代方案 1

假设不打算进行一般向上转换,则必须定义一个规则来限制将输入从 float16 -> float32 向上转换,方法是使用 DType 或 UFunc 本身(或两者的组合)上的通用逻辑。UFunc 本身无法轻松地做到这一点,因为它无法知道注册循环的所有可能的 DType。请考虑以下两个示例

首先(应拒绝)

  • 输入:float16 * float16

  • 现有循环:float32 * float32

其次(应接受)

  • 输入:timedelta64 * int32

  • 现有循环:timedelta64 * int16

这需要以下任一操作

  1. timedelta64 以某种方式发出信号,表明如果操作涉及 int64 向上转换,则始终支持该转换。

  2. float32 * float32 循环拒绝向上转换。

实现第一种方法需要在特定上下文中发出向上转换可以接受的信号。这将需要额外的挂钩,对于复杂的 DType 来说可能并不简单。

对于第二种方法,在大多数情况下,简单的 np.common_dtype 规则适用于初始分派,但是,即使对于同构循环,这种情况也是明确的。此选项将需要添加一个函数来检查输入是否是对每个循环的有效向上转换,这似乎存在问题。在许多情况下,可以提供默认值(同构签名)。

替代方案 2

另一种“提升”步骤是确保在首先找到正确的输出 DType 后,输出 DType 与循环匹配。如果知道输出 DType,则找到安全循环变得很容易。在大多数情况下,这都有效,正确的输出 dtype 只是

np.common_dtype(*input_DTypes)

或某些固定 DType(例如,逻辑函数的 Bool)。

但是,它在上述 timedelta64 * int32 案例中失败,因为事先无法知道此输出的“预期”结果类型确实是 timedelta64np.common_dtype(Datetime64, Int32) 失败)。这需要一些关于 timedelta64 精度为 int64 的额外知识。由于 ufunc 可以具有任意数量的(相关)输入,因此至少需要在 Datetime64(以及所有 DType)上添加额外的 __promoted_dtypes__ 方法。

掩码 DType 显示了进一步的限制。当涉及掩码 DType 时,逻辑函数没有布尔结果,因此这需要原始 ufunc 作者在此方案中预测掩码 DType。类似地,为复数值定义的一些函数将返回实数,而其他函数将返回复数。如果原始作者没有预测复数,则对于后来添加的复数循环,提升可能不正确。

我们认为,提升器虽然允许巨大的理论复杂性,但却是最佳解决方案

  1. 提升允许动态添加新的循环。例如,可以定义一个抽象的 Unit DType,它动态创建类来包装其他现有的 DType。使用单个提升器,此 DType 可以动态包装现有的 ArrayMethod,使其能够在单个查找中找到正确的循环,而不是两个查找。

  2. 提升逻辑通常会出错在安全的一侧:除非也添加了提升器,否则无法滥用新添加的循环。

  3. 它们将仔细考虑逻辑是否正确的负担放在向 UFunc 添加新循环的程序员身上。(与替代方案 2 相比)

  4. 如果现有提升不正确,可以编写提升器来限制或细化通用规则。通常,提升规则绝不应返回不正确的提升,但如果现有提升逻辑失败或对于新添加的循环不正确,则循环可以添加新的提升器来细化逻辑。

每个循环验证是否发生向上转换的选项可能是最佳替代方案,但不包括动态添加新循环的功能。

通用提升器的主要缺点是它们允许可能非常大的复杂性。第三方库可以向 NumPy 添加不正确的提升,但是,通过添加新的不正确循环也可以做到这一点。总的来说,我们相信我们可以依靠下游项目谨慎、负责任地使用这种能力和复杂性。

用户指南#

通常,必须非常小心地向 UFunc 添加提升器。提升器绝不应影响其他数据类型可以合理定义的循环。定义假设的 erf(UnitFloat16) 循环不应导致 erf(float16)。通常,提升器应满足以下要求

  • 在定义新的提升规则时要保守。不正确的结果比意外错误危险得多。

  • 添加的(抽象)DType 之一通常应与项目定义的 DType(或 DType 系列)具体匹配。切勿添加超出正常通用 DType 规则的提升规则!如果您编写了 int24 dtype,则添加 int16 + uint16 -> int24 的循环是不合理的。此操作的结果先前已定义为 int32,并且将在此假设下使用。

  • 提升器(或循环)绝不应影响现有循环结果。这包括添加更快速但精度较低的循环/提升器以替换现有的循环/提升器。

  • 尝试为所有与提升(和转换)相关的逻辑保持清晰的线性层次结构。NumPy 本身打破了整数和浮点数的此逻辑(它们不是严格线性的,因为 int64 无法提升到 float32)。

  • 任何项目都可以添加循环和提升器,这些项目可能是

    • 定义 ufunc 的项目

    • 定义 DType 的项目

    • 第三方项目

    尝试找出哪个项目最适合添加循环。如果定义 ufunc 的项目和定义 DType 的项目都没有添加循环,则可能会出现多个定义(被拒绝)的问题,并且应注意循环行为始终比错误更可取。

在某些情况下,这些规则的例外情况可能是有意义的,但是,总的来说,我们要求您格外小心,如有疑问,请创建一个新的 UFunc。这清楚地通知用户不同的规则。如有疑问,请在 NumPy 邮件列表或问题跟踪器中提问!

实现#

此 NEP 的实现将需要对当前 ufunc 机制(以及转换)进行大规模重构和重组。

不幸的是,实现将需要大量维护 UFunc 机制,因为必须修改实际的 UFunc 循环调用以及初始分派步骤。

通常,正确的 ArrayMethod(包括提升器返回的那些)将被缓存(或存储)在哈希表中,以便高效查找。

讨论#

存在大量可能的实现,在各个地方进行了许多讨论,以及最初的想法和设计文档。这些在 NEP 40 的讨论中列出,为了简洁起见,这里不再重复。

github 问题 12518“ufunc 内部循环签名的调用约定应该是什么?” 中可以找到涉及许多这些要点并指向类似解决方案的长时间讨论。

参考文献#

有关更多讨论和参考文献,请参见 NEP 40 和 41。