广义通用函数 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 还将包含核心维度的尺寸,从第二个元素开始。每个唯一的维度名称只提供一个尺寸,并且根据维度名称在签名中第一次出现的位置给出尺寸。

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 适用的任何值。

警告

该函数决不能更改输入时不是 -1 的 core_dim_sizes 中的值。更改不是 -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;
}