NEP 5 — 广义通用函数#

状态:

最终

正如 https://scipy.org.cn/scipy/numpy/wiki/GeneralLoopingFunctions 中所解释的，普遍需要不仅遍历标量上的函数，还需要遍历向量（或数组）上的函数。我们建议通过推广通用函数 (ufunc) 来实现这一概念，并提供一个 C 实现，该实现向 NumPy 代码库添加了约 500 行代码。在当前的（专门的）ufunc 中，基本函数仅限于逐元素操作，而广义版本支持“子数组”与“子数组”的操作。Perl 向量库 PDL 提供了类似的功能，其术语在以下内容中被重复使用。

每个广义 ufunc 都有一些与其关联的信息，这些信息声明输入的“核心”维度是什么，以及输出的相应维度（逐元素 ufunc 具有零核心维度）。所有参数的核心维度的列表称为 ufunc 的“签名”。例如，ufunc numpy.add 的签名为 (),()->()，定义了两个标量输入和一个标量输出。

另一个示例是（参见 GeneralLoopingFunctions 页面）函数 inner1d(a,b)，其签名为 (i),(i)->()。这将沿每个输入的最后一个轴应用内积，但会保持其余索引不变。例如，如果 a 的形状为 (3,5,N) 且 b 的形状为 (5,N)，则这将返回形状为 (3,5) 的输出。底层基本函数被调用 3*5 次。在签名中，我们为每个输入指定一个核心维度 (i)，为输出指定零核心维度 ()，因为它接受两个一维数组并返回一个标量。通过使用相同的名称 i，我们指定两个对应的维度应该具有相同的大小（或者其中一个的大小为 1，并将被广播）。

核心维度以外的维度称为“循环”维度。在上面的示例中，这对应于 (3,5)。

通常的 NumPy“广播”规则适用，其中签名确定每个输入/输出对象的维度如何被分割成核心维度和循环维度。

  1. 如果输入数组的维度小于相应的核心维度数，则会在其形状前添加 1。

  2. 核心维度将从所有输入中移除，剩余维度将被广播；定义循环维度。

  3. 输出由循环维度加上输出核心维度给出。

定义#

基本函数

每个 ufunc 都包含一个基本函数，该函数对数组参数的最小部分执行最基本的操作（例如，将两个数字相加是将两个数组相加的最基本操作）。ufunc 对数组的不同部分多次应用基本函数。基本函数的输入/输出可以是向量；例如，inner1d 的基本函数接受两个向量作为输入。

签名

签名是一个字符串，描述了 ufunc 的基本函数的输入/输出维度。有关更多详细信息，请参见下面的部分。

核心维度

基本函数的每个输入/输出的维度由其核心维度定义（零核心维度对应于标量输入/输出）。核心维度映射到输入/输出数组的最后维度。

维度名称

维度名称表示签名中的核心维度。不同的维度可以共享一个名称，这表示它们的大小相同（或可广播）。

维度索引

维度索引是一个整数，表示维度名称。它根据每个名称在签名中第一次出现时的顺序枚举维度名称。

签名的详细信息#

签名定义输入和输出变量的“核心”维度，从而也定义维度的收缩。签名由以下格式的字符串表示：

  • 每个输入或输出数组的核心维度由括号中的维度名称列表表示，(i_1,...,i_N)；标量输入/输出用 () 表示。可以使用任何有效的 Python 变量名来代替 i_1、i_2 等。

  • 不同参数的维度列表用 "," 分隔。输入/输出参数用 "->" 分隔。

  • 如果在多个位置使用相同的维度名称，则强制对应的维度具有相同的大小（或可广播的大小）。

签名的正式语法如下：

<Signature>            ::= <Input arguments> "->" <Output arguments>
<Input arguments>      ::= <Argument list>
<Output arguments>     ::= <Argument list>
<Argument list>        ::= nil | <Argument> | <Argument> "," <Argument list>
<Argument>             ::= "(" <Core dimension list> ")"
<Core dimension list>  ::= nil | <Dimension name> |
                           <Dimension name> "," <Core dimension list>
<Dimension name>       ::= valid Python variable name

注释

  1. 所有引号均为清晰起见。

  2. 共享相同名称的核心维度必须是可广播的，就像我们上面示例中的两个 i 一样。每个维度名称通常对应于基本函数实现中的一个循环级别。

  3. 空格将被忽略。

以下是一些签名的示例：

add

(),()->()

inner1d

(i),(i)->()

sum1d

(i)->()

dot2d

(m,n),(n,p)->(m,p)

矩阵乘法

outer_inner

(i,t),(j,t)->(i,j)

对最后维度进行内积，对倒数第二个维度进行外积，对其余维度进行循环/广播。

实现基本函数的 C-API#

当前接口保持不变，并且可以使用PyUFunc_FromFuncAndData来实现（专门的）ufunc，它由标量基本函数组成。

可以使用PyUFunc_FromFuncAndDataAndSignature来声明更通用的ufunc。参数列表与PyUFunc_FromFuncAndData相同，增加了一个参数用于指定签名作为C字符串。

此外，回调函数与以前相同，类型为void (*foo)(char **args, intp *dimensions, intp *steps, void *func)。调用时，args是一个长度为nargs的列表，包含所有输入/输出参数的数据。对于标量基本函数，steps的长度也为nargs，表示参数使用的步长。dimensions是指向单个整数的指针，定义了要循环遍历的轴的大小。

对于非平凡的签名，dimensions还将包含核心维度的尺寸，从第二个元素开始。每个唯一维度名称只提供一个尺寸，并且根据维度名称在签名中第一次出现的位置给出尺寸。

steps的前nargs个元素与标量ufunc相同。接下来的元素按顺序包含所有参数的所有核心维度的步长。

例如，考虑一个签名为(i,j),(i)->()的ufunc。在这种情况下，args将包含三个指向输入/输出数组a、b、c数据的指针。此外，dimensions将为[N, I, J]，以定义循环的大小N以及核心维度i和j的大小I和J。最后，steps将为[a_N, b_N, c_N, a_i, a_j, b_i]，包含所有必要的步长。