广义通用函数 API#

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

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

可以使用 PyUFunc_FromFuncAndDataAndSignature 来声明一个更通用的 ufunc。参数列表与 PyUFunc_FromFuncAndData 相同,但有一个额外的参数指定签名作为 C 字符串。

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

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

nargssteps 元素与标量 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],包含所有必要的步长。

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

类型为 PyUFunc_ProcessCoreDimsFunc 的可选函数,存储在 ufunc 的 process_core_dims_func 属性上,为 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;
}