广义通用函数 API#

通常需要对标量函数以及向量(或数组)函数进行循环。NumPy 通过对通用函数(ufunc)进行泛化来实现这一概念。在常规 ufunc 中,基本函数仅限于逐元素操作,而泛化版本(gufunc)支持“子数组对子数组”的操作。Perl 向量库 PDL 提供了类似的功能,其术语在下文中被重新使用。

每个广义 ufunc 都附带信息,说明输入的核心维度以及输出的相应维度(逐元素的 ufunc 具有零核心维度)。所有参数的核心维度列表称为 ufunc 的“签名”。例如,ufunc numpy.add 的签名是 (),()->(),定义了两个标量输入和一个标量输出。

另一个例子是函数 inner1d(a, b),其签名为 (i),(i)->()。这会在每个输入的最后一个轴上应用内积,但保留其余索引不变。例如,当 a 的形状为 (3, 5, N)b 的形状为 (5, N) 时,这将返回形状为 (3,5) 的输出。底层基本函数被调用 3 * 5 次。在签名中,我们为每个输入指定一个核心维度 (i),为输出指定零核心维度 (),因为它接受两个一维数组并返回一个标量。通过使用相同的名称 i,我们指定两个相应的维度应该具有相同的大小。

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

签名决定了每个输入/输出数组的维度如何划分为核心维度和循环维度

  1. 签名中的每个维度都与相应传入数组的一个维度匹配,从形状元组的末尾开始。这些是核心维度,它们必须存在于数组中,否则会引发错误。

  2. 签名中分配给相同标签的核心维度(例如 inner1d(i),(i)->() 中的 i)必须具有完全匹配的大小,不执行广播。

  3. 核心维度从所有输入中移除,剩余维度一起广播,定义循环维度。

  4. 每个输出的形状由循环维度加上输出的核心维度决定

通常,输出中所有核心维度的大小将由输入数组中具有相同标签的核心维度的大小决定。这并非强制要求,并且可以定义一个签名,其中某个标签首次出现在输出中,尽管在调用此类函数时必须采取一些预防措施。一个例子是函数 euclidean_pdist(a),其签名为 (n,d)->(p),它给定一个由 nd 维向量组成的数组,计算它们之间所有唯一的成对欧氏距离。因此,输出维度 p 必须等于 n * (n - 1) / 2,但默认情况下,调用者有责任传入正确大小的输出数组。如果输出的核心维度大小无法从传入的输入或输出数组中确定,则会引发错误。这可以通过定义一个 PyUFunc_ProcessCoreDimsFunc 函数并将其分配给 PyUFuncObject 结构的 proces_core_dims_func 字段来改变。更多详情请参见下文。

注意:在 NumPy 1.10.0 之前,检查不那么严格:必要的缺失核心维度通过在形状前添加 1 来创建,具有相同标签的核心维度会被广播,未确定的维度则创建为大小为 1。

定义#

基本函数

每个 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 | <Core dimension> |
                           <Core dimension> "," <Core dimension list>
<Core dimension>       ::= <Dimension name> <Dimension modifier>
<Dimension name>       ::= valid Python variable name | valid integer
<Dimension modifier>   ::= nil | "?"

注意事项

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

  2. 共享相同名称的未修改核心维度必须具有相同的大小。每个维度名称通常对应于基本函数实现中的一个循环级别。

  3. 空格被忽略。

  4. 将整数作为维度名称会冻结该维度为该值。

  5. 如果名称后缀带有“?”修饰符,则该维度仅当其存在于共享它的所有输入和输出上时才是一个核心维度;否则它将被忽略(并替换为基本函数中大小为 1 的维度)。

以下是一些签名的例子:

名称

签名

常见用法

add

(),()->()

二元 ufunc

sum1d

(i)->()

归约

inner1d

(i),(i)->()

向量-向量乘法

matmat

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

矩阵乘法

vecmat

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

向量-矩阵乘法

matvec

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

矩阵-向量乘法

matmul

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

以上四种的组合

outer_inner

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

在最后一个维度上进行内积,在倒数第二个维度上进行外积,并在其余维度上进行循环/广播。

cross1d

(3),(3)->(3)

叉积,其中最后一个维度被冻结且必须为 3

最后一个是冻结核心维度的实例,可用于提高 ufunc 性能

实现基本函数的 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 将包含三个指向输入/输出数组 abc 数据的指针。此外,dimensions 将是 [N, I, J],用于定义循环的 N 大小以及核心维度 ij 的大小 IJ。最后,steps 将是 [a_N, b_N, c_N, a_i, a_j, b_i],包含所有必要的步长。

自定义核心维度大小处理#

存储在 ufunc 的 process_core_dims_func 属性上的可选函数 PyUFunc_ProcessCoreDimsFunc 类型,为 ufunc 的作者提供了一个“钩子”,用于处理传递给 ufunc 的数组的核心维度。这个“钩子”的两个主要用途是:

  • 检查 ufunc 所要求的核心维度约束是否满足(如果不满足则设置异常)。

  • 计算未由输入数组确定的任何输出核心维度的输出形状。

作为第一个用法的例子,考虑签名为 (n)->(2) 的广义 ufunc minmax,它同时计算序列的最小值和最大值。它应该要求 n > 0,因为长度为 0 的序列的最小值和最大值没有意义。在这种情况下,ufunc 作者可能会这样定义该函数:

int minmax_process_core_dims(PyUFuncObject ufunc,
                             npy_intp *core_dim_sizes)
{
    npy_intp n = core_dim_sizes[0];
    if (n == 0) {
        PyExc_SetString("minmax requires the core dimension "
                        "to be at least 1.");
        return -1;
    }
    return 0;
}

在这种情况下,数组 core_dim_sizes 的长度将为 2。数组中的第二个值将始终为 2,因此函数无需检查它。核心维度 n 存储在第一个元素中。如果函数发现 n 为 0,则会设置异常并返回 -1。

“钩子”的第二个用途是,当调用者未提供输出数组且输出的一个或多个核心维度不是输入核心维度时,计算输出数组的大小。如果 ufunc 在 process_core_dims_func 属性上没有定义函数,则未指定的输出核心维度大小将导致引发异常。借助 process_core_dims_func 提供的“钩子”,ufunc 的作者可以将输出大小设置为适合 ufunc 的任何值。

在传递给“钩子”函数的数组中,未由输入确定的核心维度在 core_dim_sizes 数组中以值 -1 表示。函数可以根据输入数组中出现的核心维度,将 -1 替换为适合 ufunc 的任何值。

警告

函数绝不能更改 core_dim_sizes 中输入时非 -1 的值。更改非 -1 的值通常会导致 ufunc 的输出不正确,并可能导致 Python 解释器崩溃。

例如,考虑广义 ufunc conv1d,其基本函数分别计算两个长度为 mn 的一维数组 xy 的“完全”卷积。此卷积的输出长度为 m + n - 1。为了将其实现为广义 ufunc,签名设置为 (m),(n)->(p),并且在“钩子”函数中,如果发现核心维度 p 为 -1,则将其替换为 m + n - 1。如果 p 不是 -1,则必须验证给定值是否等于 m + n - 1。如果不等于,函数必须设置异常并返回 -1。为了得到有意义的结果,该操作还要求 m + n 至少为 1,即两个输入都不能长度为 0。

示例代码如下:

int conv1d_process_core_dims(PyUFuncObject *ufunc,
                             npy_intp *core_dim_sizes)
{
    // core_dim_sizes will hold the core dimensions [m, n, p].
    // p will be -1 if the caller did not provide the out argument.
    npy_intp m = core_dim_sizes[0];
    npy_intp n = core_dim_sizes[1];
    npy_intp p = core_dim_sizes[2];
    npy_intp required_p = m + n - 1;

    if (m == 0 && n == 0) {
        // Disallow both inputs having length 0.
        PyErr_SetString(PyExc_ValueError,
            "conv1d: both inputs have core dimension 0; the function "
            "requires that at least one input has size greater than 0.");
        return -1;
    }
    if (p == -1) {
        // Output array was not given in the call of the ufunc.
        // Set the correct output size here.
        core_dim_sizes[2] = required_p;
        return 0;
    }
    // An output array *was* given.  Validate its core dimension.
    if (p != required_p) {
        PyErr_Format(PyExc_ValueError,
                "conv1d: the core dimension p of the out parameter "
                "does not equal m + n - 1, where m and n are the "
                "core dimensions of the inputs x and y; got m=%zd "
                "and n=%zd so p must be %zd, but got p=%zd.",
                m, n, required_p, p);
        return -1;
    }
    return 0;
}