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 版本添加对 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 标准中的一些关键设计规则（例如，从输入 dtype 可预测的输出 dtype，没有具有通过关键字控制的不同返回数量的多态 API）也将应用于 NumPy 函数（这些函数不是数组 API 标准的一部分），因为这些设计规则现在被普遍认为是最佳实践。这两个设计规则尤其使 Numba 和其他 JIT 编译器更容易支持 NumPy 或与 NumPy 兼容的 API。我们将注意到，使现有的参数仅限位置和仅限关键字对于将来添加到 NumPy 的函数来说是一个好主意，但不会对现有函数这样做，因为每个此类更改都是向后兼容性中断，并且对于编写可在支持该标准的库之间移植的代码来说并非必要。现在将这些设计规则应用于主命名空间中的所有函数的另一个原因是，它使得处理 NumPy 中已经存在的新函数的潜在标准化变得容易得多——否则，由于需要向后兼容性，这些函数可能会被阻止或被迫使用替代函数名称。

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

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

好处

  • 它将使使用数组的库（例如 SciPy 和 scikit-learn，以及堆栈中较小的库）能够实现对多个数组库的支持，

  • 这将解决其他数组库在选择实现哪个 API 时遇到的“必须在 NumPy API 和数组 API 标准之间做出选择”的问题，

  • 使 CuPy、JAX、PyTorch、Dask、Numba 等库和编译器更容易匹配或支持 NumPy，方法是提供更清晰和精简的 API 接口，并解决由于 NumPy 语义导致 JIT 编译器难以支持的一些差异，

  • 一些独立于标准的新功能：添加 matrix_transpose 和 ndarray.mT，添加 vecdot，引入 matrix_norm/vector_norm（它们可以设置为 gufuncs，vecdot 已经有一个 PR 将其设置为 gufunc），

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

  • 数组 API 标准往往比 NumPy 本身具有更一致的行为（在两者之间存在差异的情况下，例如参见标准中的 线性代数设计原则 和 数据相关输出形状页面），

成本

  • 一些向后兼容性中断（大多是次要的，请参见下面的“向后兼容性”部分），

  • 主命名空间增加了大约 20 个别名（例如，acos 及其同类用 C99 名称作为 arccos 及其同类的别名）。

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

此 NEP 的范围包括

  • 为了支持 2022.12 版本的数组 API 标准，对 NumPy 的 Python API 进行的更改，包括主命名空间以及 numpy.linalg 和 numpy.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])，在某些情况下，像 unique_* 这样的新 API，在 GPU 上更容易实现，并且由于行为更可预测，对于 JIT 编译器也更容易实现。

向后兼容性#

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

  1. 在某些 NumPy 目前允许更灵活行为的地方引发错误以保证一致性/严格性，

  2. 一些逐元素函数和约简的返回数组的 dtype，

  3. 一些容差关键字的数值行为，

  4. 移动到 numpy.linalg 并支持堆叠/批处理的函数，

  5. asarray 和 array 中 copy 关键字的语义，

  6. 对 numpy.fft 功能的更改。

引发错误以保证一致性/严格性包括:

  1. 使 .T 在维度 >2 时出错，

  2. 使 cross 在大小为 2 的向量上出错（仅支持大小为 3 的向量），

  3. 使 solve 在输入不明确时出错（仅当 x2.ndim == 1 时才接受 x2 作为向量），

  4. outer 在维度 >1 的输入上引发错误而不是展平，

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

**一些逐元素函数和约简的返回数组的 dtype** 包括需要保留 dtype 的函数：ceil、floor 和 trunc 将开始返回具有相同整数 dtype 的数组，如果输入具有整数 dtype。

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

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

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

  • matrix_rank 的 tol 关键字更改为 rtol，并具有不同的解释。此外，matrix_rank 将不再支持一维数组输入，

对这些容差更改发出 FutureWarning 看起来并不合理；对于绝大多数用户来说，这些警告都是无意义的，并且它会迫使用户硬编码容差值以避免警告。数值结果的更改原则上是不可取的，因此虽然我们预计影响很小，但在主要版本中这样做会更好。

我们预计此类更改的影响中等。这是唯一一类不会导致清晰的异常或警告的更改，因此如果它确实重要（例如，下游测试开始失败或用户注意到行为发生变化），则可能需要用户付出更多努力来追踪问题。这种情况应该很少发生——在实现此更改的 PR 合并一个月后（参见 gh-25437），到目前为止报告的影响是 AstroPy 中的单个测试失败。

已移至 numpy.linalg 并支持堆叠/批处理的函数包括 diagonal 和 trace 函数。它们是标准库中 linalg 子模块的一部分，而不是主命名空间。因此，它们将在 numpy.linalg 中引入。它们将作用于最后两个轴，而不是前两个轴。这样做是为了保持一致性，因为现在其他 NumPy 函数也是这样工作的，并且为了支持“堆叠”（或其他库中更常用的术语“批处理”）。因此，linalg 和主命名空间中同名函数的行为会有所不同。从技术上讲，这不是破坏性更改，但由于同名函数的行为不同，可能会造成混淆。我们可能会弃用 np.trace 和 np.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 的实现和 PocketFFT 的 C++ 版本的移植完成的（gh-25711）。

对 numpy.fft 的另一个较小的向后不兼容的更改是，通过不允许 s 中的 None 值以及要求如果使用 s，则必须指定 axes 来简化 n 维变换中 s 和 axes 参数的行为，使其更容易理解（参见 gh-25495）。

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

适应更改和工具支持#

数组 API 的某些部分已作为 NumPy 2.0 的通用 Python API 清理工作的一部分实现（参见 NEP 52），例如

  • 建立一种命名 inf 和 nan 的方法，该方法与数组 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 迁移和在多个版本中的弃用。

详细描述#

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

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 按预期工作，从输入列表创建一个一维 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 关键字语义#

asarray 和 array 中的 copy 关键字现在将支持 True/False/None，并具有新的含义

  • True - 始终进行复制。

  • False - 永不复制。如果需要复制，则会引发 ValueError。

  • None - 只有在必要时才会进行复制（以前是 False）。

astype 中的 copy 关键字将保留其当前含义，因为在请求转换为不同的 dtype 时“永不复制”并没有多大意义。

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

新的函数名称别名#

在 NumPy 2.0 的 Python API 清理工作中（参见 NEP 52 — Python API cleanup for NumPy 2.0），我们付出了很多努力来删除别名。因此，引入新的别名必须有充分的理由。在这种情况下，需要它来匹配其他库。添加的主要别名集是用于三角函数的，其中数组 API 标准选择遵循 C99 和其他库，使用 acos、asin 等，而不是 arccos、arcsin 等。NumPy 通常也遵循 C99；目前尚不清楚多年前为什么做出这种命名选择。

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

  • 三角函数：acos、acosh、asin、asinh、atan、atanh、atan2

  • 位运算函数：bitwise_left_shift、bitwise_invert、bitwise_right_shift

  • 其他函数：concat、permute_dims、pow

  • 在 numpy.linalg 中：tensordot、matmul

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

语义重叠的新关键字#

与函数名别名类似，还有一些新的关键字与现有关键字重叠。

  • 用于`std和`var的`correction关键字（与`ddof重叠）

  • 用于`sort和`argsort的`stable关键字（与`kind重叠）

使用`correction名称是为了清晰起见（“自由度增量”不容易理解）。`stable与`kind互补，`kind已经具有`'stable'作为选项（不过，单独的关键字可能更容易发现，因此无论如何都是不错的选择），允许库保留更改/改进稳定和不稳定排序算法的权利。

新的`unique_*函数#

在数组API中，`unique函数及其影响返回元组基数的`return_index、`return_inverse和`return_counts参数被四个相应的函数取代：`unique_all、`unique_counts、`unique_inverse和`unique_values。这些新函数避免了多态性，而多态性往往是JIT编译器和静态类型检查的问题。因此，使用这些函数有助于Numba之类的工具以及Mypy之类的静态类型检查器的用户。

添加`np.bool#

曾经存在于NumPy但已被移除的别名之一是`np.bool。为了符合数组API，它被重新引入，但含义不同，因为它现在指向NumPy的bool，而不是Python的内置类型。这个改变是个好主意，我们本来也计划这样做，因为`bool比`bool_更简洁。但是，如果它不是数组API标准的一部分，我们可能不会将其重新引入2.0版本中。

未采用的标准部分#

标准规定了一些我们建议不遵循的内容（至少目前如此）。这些是：

  1. 当`dtype=None时，要求`sum和`prod始终将低精度浮点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数组的鸭子类型，标准允许这样做（它不强制只有一个数组类型，也不包含`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，而是保留在当前位置，因为它仍然有用——这是测试下游代码是否真正可以在不同数组库之间移植的最佳方法。这是一个非常合理的替代方案，但是稍微更倾向于将该模块转换成一个独立的包。

讨论#

参考文献和脚注#