NEP 56 — NumPy 主命名空间中的数组 API 标准支持#

作者:

Ralf Gommers <ralf.gommers@gmail.com>

作者:

Mateusz Sokół <msokol@quansight.com>

作者:

Nathan Goldbaum <ngoldbaum@quansight.com>

状态:

最终

替换:

NEP 30 — NumPy 数组的鸭子类型 - 实现, NEP 31 — NumPy API 的上下文局部和全局覆盖, NEP 37 — NumPy 类模块的调度协议, NEP 47 — 采用数组 API 标准

类型:

标准跟踪

创建时间:

2023-12-19

决议:

https://mail.python.org/archives/list/[email protected]/message/Z6AA5CL47NHBNEPTFWYOTSUVSRDGHYPN/

摘要#

此 NEP 提案在 NumPy 2.0 版本的 NumPy 主命名空间中添加对 2022.12 版本数组 API 标准的几乎完全支持。

在主命名空间中采用它有许多优点;最重要的是,对于依赖 NumPy 并希望开始支持其他数组库的库来说,SciPy 和 scikit-learn 是两个已经沿着这条路前进的著名库。在主命名空间中支持数组 API 标准的必要性源于从这些库以及使用不同数组对象的实验性 numpy.array_api 实现中学到的经验教训。其他数组库、像 Numba 这样的 JIT 编译器以及可能更容易在不同的数组库之间切换的最终用户也将从中受益。

动机和范围#

注意

此 NEP 中提出的主要更改在 2023 年 4 月的 NumPy 2.0 开发者会议中进行了介绍(有关该会议的演示文稿,请参见 此处),并在那里获得了认可。NumPy 2.0 的大多数实施工作已经合并。对于其余部分,PR 已准备好 - 这些主要是特定于数组 API 支持的项目,如果没有该上下文,我们可能不会考虑将其包含在 NumPy 中。此 NEP 将更详细地介绍这些 API 和 PR。

NEP 47 — 采用数组 API 标准 包含将数组 API 支持添加到 NumPy 的动机。此 NEP 扩展并取代了 NEP 47。NEP 47 旨在创建一个单独的 numpy.array_api 子模块而不是主命名空间的主要原因是,转换规则差异太大。随着基于值的转换被移除(NEP 50 — Python 标量的提升规则),这将在 NumPy 2.0 中得到解决。让 NumPy 成为数组 API 标准的超集将极大地提高代码可移植性,以便移植到其他库(CuPy、JAX、PyTorch 等),从而解决 2020 年 NumPy 用户调查[4] 中的顶级用户请求之一(GPU 支持)。请参见 numpy.array_api API 文档 (1.26.x),以概述它与主命名空间之间的差异(注意,这些“严格性”不适用)。

numpy.array_api 的体验,它仍然被标记为实验性,表明单独的严格实现和单独的数组对象主要适合测试目的,而不适合在向下游库中的常规使用。在主命名空间中提供支持可以解决此问题。因此,此 NEP 取代了 NEP 47。 numpy.array_api 模块将移至一个独立的包中,以方便更轻松地更新,而这些更新与 NumPy 版本周期无关。

数组 API 标准中的一些关键设计规则(例如,从输入数据类型预测输出数据类型,没有通过关键字控制不同数量返回值的多态 API)也将应用于 NumPy 函数,这些函数不是数组 API 标准的一部分,因为现在人们认识到这些设计规则通常是最佳实践。这两个设计规则尤其使 Numba 和其他 JIT 编译器更容易支持 NumPy 或 NumPy 兼容的 API。我们要注意,对于将来添加到 NumPy 的函数来说,使现有的参数成为位置限定和关键字限定是一个好主意,但对于现有的函数来说不会这样做,因为每个这样的更改都是一个向后兼容性中断,对于编写跨库支持该标准的代码来说,这是没有必要的。现在将这些设计规则应用于主命名空间中的所有函数的另一个原因是,这样处理现有 NumPy 中新函数的潜在标准化将变得容易得多 - 否则,这些函数可能会被阻止或被迫使用备用函数名称,因为需要向后兼容性。

添加到主命名空间的新函数必须与 NumPy 的其余部分很好地集成。因此,例如,它们应该按照预期遵循广播和其他规则,并与 NumPy 的所有数据类型一起使用,而不仅仅是标准中的数据类型。对于向后不兼容的更改也是如此(例如,线性代数函数需要以相同的方式支持批处理,并将最后两个轴视为矩阵)。因此,NumPy 应该变得更加一致,而不是更不一致。

以下是我们认为的一整套拟议更改的主要预期好处和成本。

好处

  • 它将使使用数组的库(如 SciPy 和 scikit-learn,以及更高层级的小型库)能够实现对多个数组库的支持。

  • 它将消除其他数组库在选择要实现的 API 时遇到的“必须在 NumPy API 和数组 API 标准之间做出选择”的问题。

  • CuPy、JAX、PyTorch、Dask、Numba 和其他此类库和编译器更容易匹配或支持 NumPy,方法是提供更明确和更小的 API 表面来作为目标,以及通过解决由于难以在 JIT 编译器中支持的 Numpy 语义而导致的一些差异。

  • 一些具有独立于标准的好处的新功能:添加 matrix_transposendarray.mT,添加 vecdot,引入 matrix_norm/vector_norm(它们可以被制作成 gufuncs,vecdot 已经有一个 PR 使其成为一个)。

  • NumPy 和其他数组库的 API 之间的更密切对应关系将降低最终用户在从一个数组库切换到另一个数组库时的学习曲线。

  • 数组 API 标准倾向于具有比 NumPy 本身更一致的行为(在两者的行为存在差异的情况下,例如,请参见标准中的 线性代数设计原则数据相关输出形状页面)。

成本

  • 一些向后兼容性中断(大多数是轻微的,请参见后面的“向后兼容性”部分)。

  • 使用大约 20 个别名扩展主命名空间的大小(例如,acos & co. 与 C99 名称别名 arccos & co.)。

总的来说,我们认为好处远远大于成本 - 并且是永久性的,而成本在很大程度上是暂时的。特别是,对希望与 NumPy 兼容的数组库和编译器的好处非常大。因此,整个 PyData(或科学 Python)生态系统的长期效益也很大 - 因为下游库能够更容易地支持多个数组库。所需的重大更改数量相当有限,这些更改的影响似乎很小。并非没有痛苦,但我们认为影响比 NumPy 2.0 中其他重大更改的影响要小,并且是值得付出的代价。

此 NEP 的范围包括

  • 对 NumPy 的 Python API 进行更改以支持 2022.12 版本的数组 API 标准,在主命名空间以及 numpy.linalgnumpy.fft 中。

  • 更改现有 NumPy 函数的行为,这些函数不在(或尚未在)数组 API 标准中,以与标准的关键设计原则保持一致。

此 NEP 的范围之外

  • 与数组 API 标准无关的 NumPy 的 Python API 的其他更改。

  • 对 NumPy 的 C API 的更改。

此 NEP 将取代以下 NEP

使用和影响#

我们考虑了几种不同类型的用户:编写数值代码的最终用户,依赖 NumPy 并希望开始支持多个数组库的下游包,以及旨在实现 NumPy 类或 NumPy 兼容 API 的其他数组库和工具。

从数组 API 支持中获益最明显的用户可能是想要开始支持 CuPy、PyTorch、JAX、Dask 或其他类似库的下游库。SciPy 和 scikit-learn 已经在朝着这个方向发展,并且在他们自己的 API 中的一小部分(该支持仍被标记为实验性)成功地支持 CuPy 数组和 PyTorch 张量。

他们使用的主要原则是用一个实用函数来替换常规的 import numpy as np,从输入数组中检索数组库命名空间。他们将其称为 xp,如果输入是 NumPy 数组,它实际上是 np 的别名,如果是 CuPy 数组,则是 cupy,如果是 PyTorch 张量,则是 torch。这个 xp 允许编写适用于所有这些库的代码——因为数组 API 标准是它们的共同点。例如,这段代码取自 scipy.cluster

def vq_py(obs, code_book, check_finite=True):
    """Python version of vq algorithm"""
    xp = array_namespace(obs, code_book)
    obs = as_xparray(obs, xp=xp, check_finite=check_finite)
    code_book = as_xparray(code_book, xp=xp, check_finite=check_finite)

    if obs.ndim != code_book.ndim:
        raise ValueError("Observation and code_book should have the same rank")

    if obs.ndim == 1:
        obs = obs[:, xp.newaxis]
        code_book = code_book[:, xp.newaxis]

    # Once `cdist` has array API support, this `xp.asarray` call can be removed
    dist = xp.asarray(cdist(obs, code_book))
    code = xp.argmin(dist, axis=1)
    min_dist = xp.min(dist, axis=1)
    return code, min_dist

它看起来很像普通的 NumPy 代码,但可以将 PyTorch 张量作为输入运行,然后返回 PyTorch 张量。当然,这只是这个故事的一部分。关于 scikit-learn [1] 和 SciPy [2] 的经验和影响(在某些情况下,性能大幅提升——LinearDiscriminantAnalysis.fit 在 GPU 上使用 PyTorch 比 NumPy 快 28 倍)提供了更完整的画面。

对于直接使用 NumPy 的最终用户来说,除了 NumPy 与他们可能想使用的其他库之间存在更少的差异之外,几乎没有变化。这缩短了他们的学习曲线,并使他们在 NumPy 和 PyTorch/JAX/CuPy 之间切换变得更容易。此外,他们应该从使用数组的库开始支持多个数组库中获益,使他们使用 Python 包堆栈进行科学计算或数据科学的体验更加无缝。

最后,对于其他数组库的作者以及 Numba 等工具来说,使 NumPy 与数组 API 标准保持一致的 API 改进也将节省他们的时间。由于行为更加可预测,设计规则 ([3]) 以及在某些情况下新的 API,例如 unique_*,在 GPU 和 JIT 编译器上更容易实现。

向后兼容性#

对向后兼容性产生影响的更改属于以下类别:

  1. 在某些地方提高错误,以确保一致性/严格性,而 NumPy 目前允许更灵活的行为。

  2. 某些逐元素函数和约简的返回值数组的 dtype。

  3. 几个容差关键字的数值行为。

  4. 已移动到 numpy.linalg 的函数和支持堆叠/批处理。

  5. asarrayarraycopy 关键字的语义。

  6. numpy.fft 功能的更改。

提高错误以确保一致性/严格性包括::

  1. 使 .T 在超过 2 维时出错。

  2. 使 cross 在大小为 2 的向量上出错(仅支持大小为 3 的向量)。

  3. 使 solve 在输入不明确时出错(仅接受 x2 作为向量,如果 x2.ndim == 1)。

  4. outer 在超过 1 维的输入上引发异常而不是扁平化。

我们预计此类更改的影响很小。

**某些逐元素函数和约简的返回值数组的 dtype** 包括需要保留 dtype 的函数:ceilfloortrunc 将开始返回具有相同整型 dtype 的数组,如果输入具有整型 dtype。

我们预计此类更改的影响很小。

**数值行为的更改** 包括:

  • pinvrtol 默认值从 1e-15 更改为一个依赖于 dtype 的默认值 None,解释为 max(M, N) * finfo(result_dtype).eps

  • matrix_ranktol 关键字更改为 rtol,但解释不同。此外,matrix_rank 将不再支持 1 维数组输入。

为这些容差更改引发 FutureWarning 似乎不合理;对于绝大多数用户来说,这些警告将是虚假警告,并且它会迫使用户硬编码容差值以避免警告。数值结果的更改原则上是不希望的,因此虽然我们预计影响很小,但在主要版本中进行更改会更好。

我们预计此类更改的影响中等。这是唯一一类不会导致明确异常或警告的更改,因此如果它确实很重要(例如,下游测试开始失败或用户注意到行为改变),可能需要用户花费更多时间来追踪问题。这种情况应该很少发生——在合并实现此更改的 PR 一个月后(参见 gh-25437),迄今为止报告的影响仅仅是 AstroPy 中的一个测试失败。

**已移动到 numpy.linalg 并支持堆叠/批处理的函数** 是 diagonaltrace 函数。它们是标准中的 linalg 子模块的一部分,而不是主命名空间。因此,它们将在 numpy.linalg 中引入。它们将对最后两个轴而不是前两个轴进行操作。这样做是为了保持一致性,因为现在其他 NumPy 函数的工作方式也是如此,并且是为了支持“堆叠”(或者在其他库中更常用的术语“批处理”)。因此,linalg 和主命名空间中相同名称的函数将有所不同。这在技术上不算破坏性更改,但可能令人困惑,因为具有相同名称的函数行为不同。我们可能会弃用 np.tracenp.diagonal 来解决这个问题,但最好不要立即弃用,以避免用户必须编写 if-2.0-else 条件代码。

我们预计此类更改的影响很小。

**asarrayarraycopy 关键字的语义** 将从“必要时复制”更改为“从不复制”。现在有三种类型的行为,而不是两种——copy=None 表示“必要时复制”。

我们预计此类更改的影响中等。如果用户因为显式地使用 copy=False 来进行复制但之前已经进行了复制而得到异常,则他们必须检查代码并确定代码的意图是旧语义还是新语义(两种可能性都大致相同),并根据需要调整代码。我们预计大多数情况下是 np.array(..., copy=False),因为直到几年前,这比 np.asarray(...) 具有更低的开销。这个问题已经解决,np.asarray(...) 是 NumPy 中的惯用用法。

**对 numpy.fft 的更改:** numpy.fft 子模块中的所有函数都需要为 32 位输入 dtype 保持精度,而不是向上转换为 float64/complex128。这是一个理想的更改,与 NumPy 作为整体的设计一致——但返回数组的较低精度或 dtype 可能影响用户对该模块中函数调用的影响。此更改是通过新的基于 gufunc 的实现以及在 (gh-25711) 中对 C++ 版本的 PocketFFT 进行供应商化来实现的。

numpy.fft 的另一个较小的向后不兼容更改是通过在 n 维变换中禁止 None 值来使 saxes 参数的行为更易于理解,并且要求如果使用 s,则必须指定 axes(参见 gh-25495)。

我们预计此类更改的影响很小。

适应更改和工具支持#

数组 API 的某些部分已经在 NumPy 2.0 的一般 Python API 清理工作中实现(参见 NEP 52),例如:

  • 建立一种命名 infnan 的方法,该方法与数组 API 兼容。

  • 删除隐晦的 dtype 名称,并为每个 dtype 建立(与数组 API 兼容的)规范名称。

迁移到与 NEP 52 兼容的代码库的所有说明都可以在 NumPy 2.0 迁移指南 中找到。

此外,为自动迁移 Python API 更改,实现了一个新的 ruff 规则。值得注意的是,新规则 NP201 仅用于遵守 NEP 52 更改,不涵盖使用作为数组 API 标准一部分的新函数,也不涵盖具有某些类型的向后不兼容更改的 API,如上所述。

为了自动迁移到与数组 API 兼容的代码库,正在实现一个新规则(参见问题 ruff#8615 和 PR ruff#8910)。

有了这两个规则,下游用户应该能够更新他们的项目(在自动化的范围内),以使其成为一个与库无关的代码库,该代码库可以从不同的数组库和设备中受益。

无法自动处理的向后不兼容更改(例如,线性代数函数的 rtol 默认值的更改)将以与 NumPy 2.0 中的任何其他向后不兼容更改相同的方式处理——通过文档、发行说明、API 迁移和在多个版本中进行弃用。

详细说明#

在本节中,我们将重点关注我们如果不存在标准并且不必考虑/担心其主要目标——编写可移植到多个数组库及其支持的功能(如 GPU 和其他硬件加速器或 JIT 编译器)的代码——就不会考虑引入 NumPy 的特定 API 添加和功能。

device 支持#

设备支持可能是最明显的例子。NumPy 是并且将保持一个仅限 CPU 的库,那么为什么要费心引入 ndarray.device 属性或几个函数中的 device= 关键字呢?这个特性纯粹是为了使编写可移植到各个库的代码变得更容易。 .device 属性将返回一个表示 CPU 的对象,该对象将被接受为 device= 关键字的输入。例如

# Should work when `xp` is `np` and `x1` a numpy array
x2 = xp.asarray([0, 1, 2, 3], dtype=xp.float64, device=x1.device)

这将像预期的那样在 NumPy 中工作,从输入列表中创建一个 1 维 numpy 数组。它也将适用于 CuPy 等库,它们可能会在 GPU 或其他支持的设备上创建一个新数组。

isdtype#

数组 API 标准引入了一个新的函数 isdtype 用于 dtype 的内省,因为 NumPy 中没有合适的替代方案。最接近的是 np.issubdtype,但是它假设一个复杂的类层次结构,其他数组库没有这种层次结构,它不是最符合人体工程学的设计 API,并且需要更大的 API 表面(np.floating 等)。 isdtype 将是内省 dtype 的新且规范的方式。它只需要 dtype 实现 __eq__,并在与来自同一库的其他 dtype 进行比较时具有预期的行为。

注意,作为 NEP 52 工作的一部分,删除了一些 dtype 别名,并记录了规范的 Python 和 C 名称。另请参见 gh-17325,其中涵盖了 NumPy 缺乏良好 API 的问题。

copy 关键字语义#

asarrayarray 中的 copy 关键字现在将支持 True / False / None,它们具有新的含义。

  • True - 始终创建副本。

  • False - 从不创建副本。如果需要副本,则会引发 ValueError

  • None - 仅在必要时创建副本(以前为 False)。

astype 中的 copy 关键字将保持其当前含义,因为在请求转换为不同数据类型时,“从不复制”没有太大意义。

语义变化仍然存在一个问题:如果对于用户代码 np.array(obj, copy=False),NumPy 可能会最终调用 obj.__array__,在这种情况下,将结果转换为 NumPy 数组是 obj.__array__ 的实现者的责任。因此,我们需要在 __array__ 中添加 copy=None 关键字,并将副本关键字值传递下去 - 注意在 __array__ 的实现者还没有新关键字时(在这种情况下会发出 DeprecationWarning,以便允许逐步过渡)不会破坏向后兼容性。

新的函数名别名#

在 NumPy 2.0 的 Python API 清理中(参见 NEP 52 — Python API cleanup for NumPy 2.0),我们花费了大量精力来移除别名。因此,引入新的别名必须有充分的理由。在这种情况下,需要它来匹配其他库。添加的主要别名集是用于三角函数,其中数组 API 标准选择遵循 C99 和其他库,使用 acosasin 等,而不是 arccosarcsin 等。NumPy 通常也遵循 C99;目前尚不清楚为什么在多年前做出了这种命名选择。

总共向主命名空间添加了 13 个别名,向 numpy.linalg 添加了 2 个别名。

  • 三角函数:acosacoshasinasinhatanatanhatan2

  • 位运算函数:bitwise_left_shiftbitwise_invertbitwise_right_shift

  • 其他函数:concatpermute_dimspow

  • numpy.linalg 中:tensordotmatmul

将来,NumPy 可以选择从其 __dir__ 中隐藏原始名称,以引导用户使用每个函数的首选拼写。

具有重叠语义的新关键字#

与函数名别名类似,有几个新关键字与现有关键字重叠。

  • correction 关键字用于 stdvar(与 ddof 重叠)

  • stable 关键字用于 sortargsort(与 kind 重叠)

correction 名称是为了清晰起见(“自由度增量”不容易理解)。stablekind 的补充,kind 已经有一个 'stable' 选项(一个单独的关键字可能更易于发现,因此无论如何都是一个不错的选择),允许库保留更改/改进稳定和不稳定排序算法的权利。

新的 unique_* 函数#

unique 函数,带有 return_indexreturn_inversereturn_counts 参数,这些参数会影响返回元组的基数,在数组 API 中被四个相应的函数替换:unique_allunique_countsunique_inverseunique_values。这些新函数避免了多态性,这往往是 JIT 编译器和静态类型的难题。因此,使用这些函数有助于 Numba 等工具以及 Mypy 等静态类型检查器的用户。

np.bool 添加#

以前存在于 NumPy 中但已被删除的别名之一是 np.bool。为了符合数组 API,它被重新引入,但含义不同,因为它现在指向 NumPy 的 bool,而不是 Python 内置的 bool。这种改变是一个好主意,我们一直计划这样做,因为 boolbool_ 更棒。但是,如果它不是数组 API 标准的一部分,我们可能还没有计划在 2.0 中重新引入该名称。

未采用的标准部分#

标准中有一些规定,我们建议遵循(至少目前不遵循)。它们是

  1. dtype=None 时,要求 sumprod 始终将低精度浮点数据类型上转换为 float64

    理由:这可能是破坏性的(例如, float32_arr - float32_arr.mean() 将产生一个 float64 数组,并导致内存使用量增加一倍)。虽然这种上转换已经对具有较低精度整数数据类型的输入执行,并且在那里似乎很有用以防止溢出,但对于浮点数据类型而言,这种要求似乎不太合理。

    array-api#731 已打开以重新考虑标准中的这一设计选择,该选择已在即将发布的 2023.12 版本中被接受。

  2. 在许多地方将函数签名设为仅位置和仅关键字。

    理由:2022.12 版本的标准说“必须”,但这已在即将发布的 2023.12 版本中被弱化为“应该”,以承认不这样做是可以的 - 毕竟数组库的用户仍然可以使用推荐的风格编写他们的代码。对于 NumPy,这些更改将是有用的,并且似乎我们可能随着时间的推移引入其中许多或全部更改(实际上 ufunc 已经符合标准),但是没有必要急于进行此更改 - 为 2.0 进行此更改将是不必要的破坏性。

  3. 要求“就地操作必须与各自的二进制(即,两个操作数,非赋值)操作具有相同的行为(包括特殊情况)”(不包括对视图的影响)。

    理由:该要求非常合理,并且可能是大多数 NumPy 用户预期的行为。但是,弃用就地运算符的非安全转换是一个难以预测影响的更改。因此,需要首先调查此事,然后如果影响足够低,则可能可以根据 NumPy 的正常向后兼容性指南弃用当前行为。

    此主题在 gh-25621 中跟踪。

注意

我们注意到,NumPy 特定的行为仍然是返回数组标量,而不是 0 维数组,在大多数情况下标准和其他数组库返回 0 维数组(例如,索引和缩减)。数组标量基本上是 0 维数组的鸭子类型,这在标准中是允许的(它不规定只有一种数组类型,也不包含 isinstance 检查或其他与数组标量不兼容的语义)。在过去的一年中,已经进行了多次讨论,讨论了从 NumPy 中删除数组标量,或者至少不再默认返回它们的可行性。但是,这将是一个巨大的工作量,在技术风险和更改的影响方面存在一些不确定性,而且还没有人着手去做。鉴于数组标量实现了很大程度上与数组兼容的接口,这似乎不是关于数组 API 标准兼容性(或总体上)的最高优先级项目。

实现#

数组 API 标准支持的跟踪问题(gh-25076)记录了实现全面支持的进度,并链接到相关的讨论。它列出了所有相关 PR(已合并和待处理),这些 PR 验证或提供数组 API 支持。

由于 NEP 52 在一定程度上与这个 NEP 融合,因此我们可以在其跟踪问题(gh-23999)上找到一些相关的实现和讨论。

作为第一个合并的 PR 之一,它包含了一个新的 CI 任务,该任务添加了 array-api-tests 测试套件。这样我们就能更好地控制每次添加的函数/别名批次,并确保实现符合数组 API 标准(参见 gh-25167)。

然后,我们继续一次合并一批,添加一个特定的 API 部分。下面列出了一些更重要的部分,包括我们在本 NEP 的前面部分中讨论过的一些部分。

替代方案#

在 NumPy 的主命名空间中实现对数组 API 标准的支持的替代方案包括

  • 一个或多个被取代的 NEP,或

  • 使 ndarray.__array_namespace__() 返回一个隐藏的命名空间(甚至是一个新的公共命名空间),其中包含兼容的函数,

  • 根本不实现对数组 API 标准的支持。

与数组 API 标准相比,被取代的 NEP 都有一些缺点,到目前为止,已经投入了大量工作来构建标准,以及其他关键库的采用。因此,这些替代方案并不吸引人。鉴于人们对这个主题的兴趣,什么也不做也不吸引人。“隐藏命名空间”选项将是对此提议的较小更改。我们更愿意不这样做,因为它会导致重复的实现保留下来,实现更复杂(例如,静态类型可能出现问题),而且仍然存在两种本质上相同 API 的变体。

从 NumPy 中删除 numpy.array_api 的一个替代方案是将其保留在当前位置,因为它仍然有用 - 它是测试下游代码是否真正可以在数组库之间移植的最佳方法。这是一个非常合理的替代方案,但是稍微偏向于将该模块变成一个独立的包。

讨论#

参考资料和脚注#