NEP 5 — 广义通用函数#

状态:

最终

正如在 https://scipy.org.cn/scipy/numpy/wiki/GeneralLoopingFunctions 所解释的那样,普遍需要不仅对标量函数进行循环,还对向量(或数组)函数进行循环。我们建议通过泛化通用函数(ufuncs)来实现这一概念,并提供一个 C 实现,该实现将为 NumPy 代码库增加约 500 行代码。在当前的(专用)ufuncs 中,基本函数仅限于逐元素操作,而广义版本支持“子数组”对“子数组”的操作。Perl 向量库 PDL 提供了类似的功能,其术语将在下文重复使用。

每个广义 ufunc 都包含与其关联的信息,这些信息说明了输入的“核心”维度以及输出的相应维度(逐元素 ufuncs 的核心维度为零)。所有参数的核心维度列表称为 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_1i_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 仍然可以用于实现(专用的)ufuncs,这些 ufuncs 由标量基本函数组成。

可以使用 PyUFunc_FromFuncAndDataAndSignature 来声明一个更通用的 ufunc。参数列表与 PyUFunc_FromFuncAndData 相同,并增加了一个额外的参数,该参数将签名指定为 C 字符串。

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

对于非平凡的签名,dimensions 还将包含核心维度的大小,从第二个条目开始。每个唯一的维度名称只提供一个大小,并且大小根据维度名称在签名中的第一次出现顺序给出。

steps 的前 nargs 个元素与标量 ufuncs 相同。后续元素按顺序包含所有参数的所有核心维度的步幅。

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