通用通用函数 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 的维度替换)。

以下是一些签名的示例

名称

signature

通用用法

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

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

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

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

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

作为第一个用例的例子,考虑广义 ufunc minmax,其签名是 (n)->(2),它同时计算序列的最小值和最大值。它应该要求 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) {
        PyErr_SetString(PyExc_ValueError,
                        "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,其基本函数计算两个一维数组 xy(长度分别为 mn)的“全”卷积。此卷积的输出长度为 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;
}