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 对象将是描述任何操作数组的函数的新首选方法。

  • Ufunc 将调度到 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#

“context” 对象类似于 Python 的 self,它被传递给所有方法。要理解为什么需要“context”对象及其内部结构,记住 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 操作多个输入数组,因此操作多个数据类型。

  • 上面的 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

UFuncs 需要与类型转换相同的信息,给出以下定义(另请参见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 的某些实例(例如,只有本机字节序)。这由 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 不需要此额外的 scratch-space,但可能需要简单的标志功能,例如实现警告(参见下面的错误和警告处理)。为了简化这一点,当 user_data 未设置时,NumPy 将传递单个零初始化的 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 调用的步骤中定义了一个统一、清晰的 API 用于步骤 3。此步骤需要分配输出数组,并且必须在准备类型转换之前发生。

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

  • -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 中,我们建议传入该结构以简化参数化数据类型的循环创建。

传入用户数据

当前的类型转换实现使用现有的 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;
}

TODO:当未设置 innerloop_data 时,传入 npy_intp 临时空间的想法似乎很方便,但我对此不确定,因为我不知道任何类似的先前技术。原则上,这个“临时空间”也可以是 context 的一部分。

重用现有循环/实现#

对于许多数据类型,上述添加额外 C 级循环的定义将足够,只需要单个步幅循环实现,如果循环与参数化数据类型一起工作,则还_必须_提供 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,它为每种数值 dtype(可以动态创建)都有子类。

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

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

替代用例

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

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

提升和调度#

NumPy ufunc 在某种意义上是多方法,因为它们同时操作(或使用)多个 DType。虽然输入(和输出)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),没有直接匹配,需要一个提升步骤。

提升和调度过程#

一般来说,ArrayMethod 是通过搜索所有输入 DType 的精确匹配来找到的。输出 dtype *不应*影响计算,但是如果多个注册的 ArrayMethod 完全匹配,则将使用输出 DType 来查找更好的匹配。这将允许当前对 np.equal 循环的区分,该循环定义了 Object, Object -> Bool(默认)和 Object, Object -> Object

最初,ArrayMethod 将仅针对*具体* DType 定义,并且由于这些 DType 不能是子类,因此可以保证精确匹配。将来,我们期望 ArrayMethod 也可以针对*抽象* DType 定义。在这种情况下,最佳匹配将如下所述找到。

提升

如果存在匹配的 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 和提升器的调度规则

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

  • 签名匹配所有输入 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)在一个有序的层次结构中。

虽然“安全”转换规则不够严格,但我们可以想象使用新的“提升”转换规则或公共 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 参与运算,则始终支持 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显示了另一个限制。当涉及掩码时,逻辑函数没有布尔结果,因此这需要原始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 issue 12518 “ufunc内部循环签名的调用约定应该是什么?”中找到。

参考文献#

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