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

作者:

Ralf Gommers <ralf.gommers@gmail.com>

作者:

Mateusz Sokół <msokol@quansight.com>

作者:

Nathan Goldbaum <ngoldbaum@quansight.com>

状态:

最终

替换:

NEP 30 — Duck typing for NumPy arrays - implementation, NEP 31 — Context-local and global overrides of the NumPy API, NEP 37 — A dispatch protocol for NumPy-like modules, NEP 47 — Adopting the array API standard

类型:

标准轨道

创建时间:

2023-12-19

解决时间:

https://mail.python.org/archives/list/numpy-discussion@python.org/message/Z6AA5CL47NHBNEPTFWYOTSUVSRDGHYPN/

摘要#

此 NEP 提议为 NumPy 2.0 版本的主命名空间添加对 Array API 标准 2022.12 版本的几乎完全支持。

在主命名空间中采用具有诸多优势;对依赖 NumPy 并希望开始支持其他数组库的库而言,这一点最为重要。SciPy 和 scikit-learn 是两个已经在这条道路上取得进展的知名库。在主命名空间中支持 Array API 标准的需求源于这些库的经验教训以及实验性的 numpy.array_api 实现(使用了不同的数组对象)。其他数组库、Numba 等 JIT 编译器以及可能更容易在不同数组库之间切换的最终用户也将受益。

动机与范围#

注意

本 NEP 中提出的主要变更已在 2023 年 4 月的 NumPy 2.0 开发者会议上展示(请参阅该会议的演示文稿 此处),并在会上获得了认可。NumPy 2.0 的大部分实现工作已经合并。对于剩余部分,PRs 已经准备就绪——这些主要是 Array API 支持特有的项目,我们可能不会在没有该背景的情况下考虑将其纳入 NumPy。本 NEP 将更详细地关注这些 API 和 PRs。

NEP 47 — Adopting the array API standard 包含了将 Array API 支持添加到 NumPy 的动机。本 NEP 扩展并取代了 NEP 47。NEP 47 旨在为 numpy.array_api 子模块而非主命名空间设定的主要原因是强制转换规则差异过大。随着基于值的强制转换的移除(NEP 50 — Promotion rules for Python scalars),这一点将在 NumPy 2.0 中得到解决。将 NumPy 作为 Array 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 发布周期无关的更新。

Array API 标准中的一些关键设计规则(例如,输出 dtypes 可从输入 dtypes 预测,没有由关键字控制的可变返回数量的多态 API)也将应用于不属于 Array API 标准的 NumPy 函数,因为这些设计规则现在被认为普遍是好的实践。特别是这两条设计规则使 Numba 和其他 JIT 编译器更容易支持 NumPy 或 NumPy 兼容的 API。我们将指出,将现有参数设为仅位置参数和仅关键字参数是将来添加到 NumPy 的函数的良好做法,但不会对现有函数这样做,因为每一次此类更改都会破坏向后兼容性,并且对于编写可在支持该标准的库之间移植的代码并非必需。现在将这些设计规则应用于主命名空间中所有函数的一个额外原因是,这样可以更容易地处理 NumPy 中已存在的新函数的潜在标准化——否则,由于需要向后兼容性,这些函数可能会被阻止或被迫使用替代函数名。

重要的是,添加到主命名空间的新函数应与其他 NumPy 部分良好集成。因此,它们应遵循广播等规则,并支持 NumPy 的所有 dtypes,而不仅仅是标准中的那些。向后不兼容的更改也是如此(例如,线性代数函数需要以相同的方式支持批处理,并将最后两个轴视为矩阵)。因此,NumPy 应该变得更加一致而不是不那么一致。

以下是我们认为提议的全部变更集的主要预期好处和成本:

好处

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

  • 它将消除其他数组库在选择实现哪个 API 时所面临的“不得不在 NumPy API 和 Array API 标准之间做出选择”的问题,

  • 通过提供更明确、更小的 API 表面积,以及解决 NumPy 语义导致 JIT 编译器难以支持的一些差异,使 CuPy、JAX、PyTorch、Dask、Numba 以及其他类似库和编译器更容易匹配或支持 NumPy,

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

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

  • Array API 标准通常比 NumPy 本身具有更一致的行为(在两者之间存在差异的情况下,请参阅标准中的 线性代数设计原则数据依赖输出形状页面),

成本

  • 一些向后兼容性中断(主要是次要的,请参阅下文的向后兼容性部分),

  • 将主命名空间的大小扩展约 20 个别名(例如,acos 等 C99 名称别名为 arccos 等)。

总的来说,我们认为好处远远超过成本——而且是永久性的,而成本在很大程度上是暂时的。特别是,对于想要实现与 NumPy 兼容性的数组库和编译器而言,其好处是显著的。因此,对于 PyData(或科学 Python)生态系统作为一个整体的长期好处——因为下游库能够更容易地支持多个数组库——也是显著的。所需的破坏性更改数量相当有限,而且这些更改的影响似乎很小。虽然不是毫无痛苦,但我们认为其影响小于 NumPy 2.0 中其他破坏性更改的影响,并且是值得付出的代价。

此 NEP 的范围包括

  • NumPy 的 Python API 更改,以支持主命名空间中以及 numpy.linalgnumpy.fft 中的 Array API 标准 2022.12 版本,

  • 现有 NumPy 函数(尚未或尚未出现在 Array API 标准中)的行为更改,以使其符合标准的关键设计原则。

此 NEP 的范围不包括

  • NumPy 的其他 Python API 更改,与 Array API 标准无关,

  • NumPy 的 C API 更改。

本 NEP 将取代以下 NEP:

用法与影响#

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

从 Array API 支持中受益最显著的用户可能是那些希望开始支持 CuPy、PyTorch、JAX、Dask 或其他类似库的下游库。SciPy 和 scikit-learn 在这方面已经取得了相当大的进展,并在其部分 API 中成功支持 CuPy 数组和 PyTorch 张量(该支持仍被标记为实验性)。

他们使用的主要原则是,他们将常规的 import numpy as np 替换为一个实用函数,用于从输入数组中检索数组库命名空间。他们称之为 xp,如果输入是 NumPy 数组,它实际上是 np 的别名,对于 CuPy 数组是 cupy,对于 PyTorch 张量是 torch。这个 xp 允许编写适用于所有这些库的代码——因为 Array 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 与 Array API 标准保持一致的 API 改进也将为他们节省时间。设计规则([3]),以及在某些情况下像 unique_* 系列这样的新 API,由于行为更可预测,因此在 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.ndim == 1 时才接受 x2 作为向量),

  4. outer 函数在输入大于 1 维时会展开而不是抛出错误,

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

某些逐元素函数和归约函数的返回数组的 dtype 包括需要保留 dtypes 的函数:ceilfloortrunc 将开始返回与输入具有相同整数 dtype 的数组(如果输入是整数 dtype)。

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

数值行为更改 包括:

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

  • 传递给 matrix_ranktol 关键字更改为 rtol,其解释不同。此外,matrix_rank 将不再支持 1-D 数组输入,

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

我们预计此类更改的影响是中等的。这是唯一一类不会导致明显异常或警告的更改,因此,如果它确实重要(例如,下游测试开始失败或用户注意到行为变化),它可能需要用户花费更多精力来查找问题。这种情况应该不频繁——在实现此更改的 PR 合并一个月后(参见 gh-25437),到目前为止报告的影响是 AstroPy 中单个测试失败。

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

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

asarray 和 array 中 copy 关键字的语义对于 copy=False 将从“需要时复制”更改为“从不复制”。现在有三种行为而不是两种——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) 中引入 PocketFFT 的 C++ 版本来实现的。

numpy.fft 的一项较小的向后不兼容的更改是,通过不允许 s 中的 None 值,并要求如果使用 s,则 axes 也必须被指定,来使 n-D 变换中 saxes 参数的行为更容易理解(参见 gh-25495)。

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

适应更改和工具支持#

Array API 的一些部分已经作为 NumPy 2.0 的通用 Python API 清理的一部分实现(参见 NEP 52),例如:

  • 建立命名 infnan 的一种且唯一的方式,以兼容 Array API。

  • 移除晦涩的 dtype 名称,并为每个 dtype 建立(兼容 Array API 的)规范名称。

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

此外,还实现了一个新的 ruff 规则,用于自动迁移 Python API 更改。值得指出的是,新规则 NP201 仅用于遵循 NEP 52 更改,而不涵盖使用 Array API 标准的新函数,也不涵盖上面讨论的某些类型的向后不兼容 API。

为了自动迁移到 Array API 兼容的代码库,正在实现一个新的规则(参见 issue 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-D numpy 数组。它也将适用于 CuPy 等库,在那里它可能在 GPU 或其他支持的设备上创建一个新数组。

isdtype#

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

请注意,作为 NEP 52 工作的一部分,一些 dtype 别名被移除,并且记录了规范的 Python 和 C 名称。另请参阅 gh-17325,其中涵盖了 NumPy 缺乏良好 API 来解决此问题的相关问题。

copy 关键字语义#

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

  • True - 始终创建一个副本。

  • False - 绝不创建副本。如果需要副本,将引发 ValueError

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

由于“要求转换为不同的 dtype 时从不复制”的意义不大,astype 中的 copy 关键字将保持其当前含义。

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

新的函数名别名#

在 NumPy 2.0 的 Python API 清理工作中(请参阅 NEP 52 — NumPy 2.0 的 Python API 清理),我们花费了大量精力来移除别名。因此,引入新别名必须有充分的理由。在此情况下,为了与其他库保持一致,需要引入新别名。最主要的别名集合是为三角函数添加的,其中数组 API 标准选择遵循 C99 和其他库,使用 acosasin 等而不是 arccosarcsin 等。NumPy 通常也遵循 C99;不完全清楚多年前为何做出此命名选择。

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

  • 三角函数:acosacoshasinasinhatanatanhatan2

  • 按位函数:bitwise_left_shiftbitwise_invertbitwise_right_shift

  • 其他函数:concatpermute_dimspow

  • numpy.linalg 中:tensordotmatmul

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

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

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

  • stdvarcorrection 关键字(与 ddof 重叠)。

  • sortargsortstable 关键字(与 kind 重叠)。

correction 名称是为了清晰(“自由度差值”不易理解)。 stable 是对 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 是一个比 bool_ 更好的名称。然而,如果不是数组 API 标准的一部分,我们可能不会为 2.0 版本安排重新引入此名称。

未采纳的标准部分#

标准中规定了一些我们建议(至少目前)遵循的内容。它们是:

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

    理由:这可能具有颠覆性(例如, float32_arr - float32_arr.mean() 将产生一个 float64 数组,内存使用量翻倍)。虽然对于具有低精度整数 dtype 的输入已经进行了这种向上转换,并且这似乎有助于防止溢出,但要求对浮点 dtype 进行此操作似乎不太合理。

    array-api#731 已被提出以重新考虑标准中的此设计选择,并且已接受在下一个标准版本中进行。

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

    理由:标准的 2022.12 版本说“必须”,但这已经在即将发布的 2023.12 版本中软化为“应该”,以认识到不这样做也没关系——数组库的用户仍然可以按照推荐的风格编写代码。对于 NumPy 来说,这些更改将是有用的,并且很可能随着时间的推移(实际上 ufunc 已经符合要求)引入许多或全部更改,但没有必要急于进行此更改——为 2.0 版本这样做将不必要地具有颠覆性。

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

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

    此主题在 gh-25621 中进行跟踪。

注意

我们注意到,一个仍然存在的 NumPy 特定的行为是,在大多数标准和其他数组库返回 0-D 数组的情况下(例如,索引和归约),仍然返回数组标量而不是 0-D 数组。数组标量基本上是 0-D 数组的“duck typing”,这被标准允许(它不强制要求只有一个数组类型,也不包含 isinstance 检查或其他不起作用的语义)。在过去一年中,关于从 NumPy 中移除数组标量或至少不再默认返回它们的可能性,已经进行了多次讨论。然而,这将是一项重大工作,并且在技术风险和更改影响方面存在一定的不确定性,并且没有人承担这项工作。鉴于数组标量实现了很大程度上兼容数组的接口,这似乎不是与数组 API 标准兼容性(或总体而言)最高优先级的项目。

实现#

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

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

最早合并的 PR 之一包含了一个新的 CI 作业,该作业添加了 array-api-tests 测试套件。通过这种方式,我们可以更好地控制每次添加的函数/别名批次,并确保实现符合数组 API 标准(请参阅 gh-25167)。

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

替代方案#

在 NumPy 中实现对数组 API 标准支持的替代方案包括:

  • 一个或多个已废弃的 NEP,或

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

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

已废弃的 NEP 与数组 API 标准相比都存在一些缺点,而且到目前为止,已经为该标准付出了大量努力——以及其他关键库的采纳。因此,这些替代方案不具吸引力。鉴于对此主题的兴趣,什么都不做也不具吸引力。“隐藏命名空间”选项将是对此提案的一个较小的更改。我们倾向于不这样做,因为它会导致重复实现并存,实现更复杂(例如,可能存在静态类型问题),并且仍然存在两种本质上相同的 API 形式。

从 NumPy 中移除 numpy.array_api 的一个替代方案是将其保留在当前位置,因为它仍然有用——它是测试下游代码是否真正可移植到不同数组库的最佳方法。这是一个非常合理的替代方案,但是,我们更倾向于将该模块提取出来并将其变成一个独立的包。

讨论#

参考文献和脚注#