NEP 20 — 通用函数签名的扩展#

作者:

Marten van Kerkwijk <mhvk@astro.utoronto.ca>

状态:

最终

类型:

标准跟踪

创建:

2018-06-10

解决:

https://mail.python.org/pipermail/numpy-discussion/2018-April/077959.html, https://mail.python.org/pipermail/numpy-discussion/2018-May/078078.html

注意

添加固定 (i) 和灵活 (ii) 维度的建议被接受,而添加可广播 (iii) 维度的建议被推迟。

摘要#

通用函数,顾名思义,是通用函数的泛化:它们作用于非标量元素。它们的签名描述了它们作用的元素的结构,其中名称连接了应相同的操作数的维度。在这里,建议扩展签名,以允许签名指示维度 (i) 具有固定大小;(ii) 可以不存在;以及 (iii) 可以广播。

详细描述#

该提案的每个部分都由特定的需求驱动 [1]

  1. 固定大小的维度。处理空间向量的代码通常明确是针对 2 或 3 维空间的(例如,来自 基本天文学标准 的代码,作者希望使用 gufuncs 包装它以用于 astropy [2])。签名应该能够指示这一点。例如,将极角转换为二维笛卡尔单位向量的函数的签名目前必须是 `()->(n)`,而没有办法指示 `n` 必须等于 2。事实上,这个签名特别令人烦恼,因为如果没有输出参数,当前的 gufunc 包装代码会失败,因为它无法确定 `n`。类似地,两个三维向量的叉积的签名必须是 `(n),(n)->(n)`,同样没有办法指示 `n` 必须等于 3。因此,这里提出了一个允许在变量名称之外给出数值值的建议。因此,角度到二维单位向量的签名将是 `()->(2)`;两个角度到三维单位向量的签名将是 `(),()->(3)`;而两个三维向量的叉积的签名将是 `(3),(3)->(3)`。

  2. 可能缺少的维度。这一部分几乎完全是由希望在 gufunc 中包装 `matmul` 驱动的。`matmul` 代表矩阵乘法,如果它只做这件事,则可以使用签名 `(m,n),(n,p)->(m,p)` 来覆盖它。但是,它有一些特殊情况,当维度缺失时,允许将任一参数视为单个向量,从而有效地将函数变为向量-矩阵、矩阵-向量或向量-向量乘法(但没有广播)。为了支持这一点,建议允许将维度名称后缀为问号以指示维度不一定必须存在。

    通过此添加,`matmul` 的签名可以表示为 `(m?,n),(n,p?)->(m?,p?)`。这表示如果,例如,第二个操作数只有一个维度,则对于基本函数来说,它将被视为输入具有核心形状 `(n, 1)`,输出具有相应的核心形状 `(m, 1)`。然而,实际输出数组删除了灵活维度,即它将具有形状 `(..., m)`。类似地,如果两个参数都只有一个维度,则输入将被表示为具有形状 `(1, n)` 和 `(n, 1)` 的基本函数,输出为 `(1, 1)`,而返回的实际输出数组将具有形状 `()`。通过这种方式,签名允许使用单个基本函数来处理四个相关但不同的签名,` (m,n),(n,p)->(m,p)`, `(n),(n,p)->(p)`, `(m,n),(n)->(m)` 和 `(n),(n)->()`。

  3. 可以广播的维度。对于某些应用,操作数之间的广播是有意义的。例如,一个 `all_equal` 函数,它比较数组中的向量,可以有签名 `(n),(n)->()`,但这会强制两个操作数都是数组,而实际上也希望检查向量的所有部分是否都是常数(可能是零)。建议允许 gufunc 的实现者指示维度可以通过在维度名称后缀 `|1` 来广播。因此,`all_equal` 的签名将变为 `(n|1),(n|1)->()`。该签名似乎更普遍地适用于“链接的 UFuncs”;例如,另一个应用可能是在推定的实现 `sumproduct` 的 Ufunc 中。

    在讨论中提出的另一个例子是加权平均,它可能看起来像 `weighted_mean(y, sigma[, axis, ...])`,返回平均值及其不确定度。使用 `(n),(n)->(),()` 的签名,将被迫始终提供与数据点一样多的 sigma,而广播允许提供所有点的单个 sigma(这对于计算平均值的误差仍然有用)。

实现#

所提出的更改已全部实现 [3], [4], [5]。这些 PR 将 ufunc 结构扩展为两个新字段,每个字段的大小等于不同维度的数量,其中 `core_dim_sizes` 保存可能固定的尺寸,`core_dim_flags` 保存指示维度是否可以缺失或可广播的标志。为了确保我们可以区分此新版本和以前的版本,一个未使用的条目 `reserved1` 被重新用作版本号。

在实现中,要注意的是,对于标记的维度,基本函数不会被视为任何不同:例如,固定维度的大小仍然传递给基本函数(但循环现在可以依赖于该大小等于签名中给出的固定大小)。

一个待决定的实现细节是,是否有可能方便地总结所有标志。这可能存储在 `core_enabled` 中(目前是一个布尔值),非零值继续表示 gufunc,但具体标志指示 gufunc 是否使用固定、灵活或可广播的维度。

以上,语法的正式定义将变为 [4]

<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" | "?"
  1. 所有引号仅为清晰起见。

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

  3. 空格将被忽略。

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

  5. 如果名称后缀为 `|1` 修饰符,则允许它与具有相同名称的其他维度广播。所有输入维度都必须共享此修饰符,而没有输出维度应该具有它。

  6. 如果名称后缀为 `?` 修饰符,则维度只有在所有共享它的输入和输出上都存在时才是核心维度;否则它将被忽略(并替换为基本函数的 1 大小的维度)。

签名示例 [4]

签名

可能的用途

(),()->()

加法

(i)->()

对最后一轴求和

(i|1),(i|1)->()

沿轴测试相等性,允许与标量进行比较

(i),(i)->()

内向量积

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

矩阵乘法

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

向量-矩阵乘法

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

矩阵-向量乘法

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

一次执行以上所有四种,除了向量不能具有循环维度(即,类似于 `matmul`)

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

三维向量的叉积

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

对最后一维求内积,对倒数第二维求外积,并对其余维度循环/广播。

向后兼容性#

一个可能的担忧是 ufunc 结构的变化。对于大多数调用 `PyUFunc_FromDataAndSignature` 的应用程序来说,这完全是透明的。此外,通过将 `reserved1` 重新用作版本号,针对旧版本 numpy 编译的代码将继续工作(尽管在导入该代码的新版本 numpy 时会收到警告),除非代码显式更改 `reserved1` 条目。

替代方案#

有人建议,而不是扩展签名,使用多重调度,以便例如,`matmul` 只是拥有它支持的多重签名,即,而不是 `(m?,n),(n,p?)->(m?,p?)`,将会有 `(m,n),(n,p)->(m,p) | (n),(n,p)->(p) | (m,n),(n)->(m) | (n),(n)->()`。这样做的缺点是开发人员现在必须确保基本函数可以处理这些不同的签名。此外,扩展很快就会变得繁琐。例如,对于 `all_equal` 签名 `(n|1),(n|1)->()`,必须有五个条目:` (n),(n)->() | (n),(1)->() | (1),(n)->() | (n),()->() | (), (n)->()`。对于类似于 `(m|1,n|1,o|1),(m|1,n|1,o|1)->()`(来自 [4] 中的 `cube_equal` 测试用例)的签名,甚至不值得写出扩展。

对于广播,有人建议使用 `^` 的替代后缀(因为广播可以被认为是增加数组的大小)。这似乎不太清楚。此外,有人怀疑它是否不应该只是一个全部或无标志。这可能是这种情况,尽管鉴于灵活维度的后缀,另一个后缀更清晰(就像实现一样)。

讨论#

这些建议在邮件列表中进行了相当长的讨论 [6], [7]。主要争论点是用例是否足够强。特别是对于冻结维度,有人认为可以在循环选择代码中添加对正确数量的检查。这似乎不太清楚,没有好处。

对于广播,有人注意到没有基本函数示例可能需要它,并质疑是否像 `all_equal` 这样的东西最好用 gufunc 完成,而不是作为 `np.equal` 上的特殊方法。一个反驳这一点的论点是,有一个实际的 PR 用于 `all_equal` [8]。另一个是,即使使用方法,也应该能够表达它们的签名(就像至少对于 `reduce` 和 `accumulate` 可能的那样)。

最后的论点是,我们使 gufuncs 太复杂了。这对于可以省略的维度来说可能是正确的,但这也有最强的用例。冻结维度的实现非常简单,它的含义也很明显。广播的能力也很简单,一旦支持灵活维度。

参考文献和脚注#