Ctrl+K

支持数组 API 标准#

注意

数组 API 标准支持目前仍处于实验阶段，并隐藏在环境变量后面。目前仅覆盖了公共 API 的一小部分。

本指南介绍了如何使用和添加对Python 数组 API 标准的支持。此标准允许用户开箱即用地将任何与数组 API 兼容的数组库与 SciPy 一起使用。

该RFC定义了 SciPy 如何实现对该标准的支持，其主要原则是“输入数组类型等于输出数组类型”。此外，该实现对允许的类数组输入进行了更严格的验证，例如拒绝 numpy 矩阵和掩码数组实例，以及具有对象类型的数组。

在下文中，与数组 API 兼容的命名空间记为 xp。

使用数组 API 标准支持#

要启用数组 API 标准支持，必须在导入 SciPy 之前设置环境变量。

export SCIPY_ARRAY_API=1

这既启用了数组 API 标准支持，也启用了对数组类参数更严格的输入验证。请注意，此环境变量旨在临时使用，作为一种进行增量更改并将它们合并到 ``main`` 中而不立即影响向后兼容性的方法。我们不打算长期保留此环境变量。

此聚类示例展示了使用 PyTorch 张量作为输入和返回值的用法。

>>> import torch
>>> from scipy.cluster.vq import vq
>>> code_book = torch.tensor([[1., 1., 1.],
...                           [2., 2., 2.]])
>>> features  = torch.tensor([[1.9, 2.3, 1.7],
...                           [1.5, 2.5, 2.2],
...                           [0.8, 0.6, 1.7]])
>>> code, dist = vq(features, code_book)
>>> code
tensor([1, 1, 0], dtype=torch.int32)
>>> dist
tensor([0.4359, 0.7348, 0.8307])

请注意，以上示例适用于 PyTorch CPU 张量。对于 GPU 张量或 CuPy 数组，vq 的预期结果是 TypeError，因为 vq 不是纯 Python 函数，因此无法在 GPU 上运行。

更严格的数组输入验证将拒绝 np.matrix 和 np.ma.MaskedArray 实例，以及具有 object 数据类型的数组。

>>> import numpy as np
>>> from scipy.cluster.vq import vq
>>> code_book = np.array([[1., 1., 1.],
...                       [2., 2., 2.]])
>>> features  = np.array([[1.9, 2.3, 1.7],
...                       [1.5, 2.5, 2.2],
...                       [0.8, 0.6, 1.7]])
>>> vq(features, code_book)
(array([1, 1, 0], dtype=int32), array([0.43588989, 0.73484692, 0.83066239]))

>>> # The above uses numpy arrays; trying to use np.matrix instances or object
>>> # arrays instead will yield an exception with `SCIPY_ARRAY_API=1`:
>>> vq(np.asmatrix(features), code_book)
...
TypeError: 'numpy.matrix' are not supported

>>> vq(np.ma.asarray(features), code_book)
...
TypeError: 'numpy.ma.MaskedArray' are not supported

>>> vq(features.astype(np.object_), code_book)
...
TypeError: object arrays are not supported

当前支持的功能#

当设置环境变量时，以下模块提供数组 API 标准支持。

Support is provided in scipy.special for the following functions: scipy.special.log_ndtr, scipy.special.ndtr, scipy.special.ndtri, scipy.special.erf, scipy.special.erfc, scipy.special.i0, scipy.special.i0e, scipy.special.i1, scipy.special.i1e, scipy.special.gammaln, scipy.special.gammainc, scipy.special.gammaincc, scipy.special.logit, and scipy.special.expit.

实现说明#

对数组 API 标准的支持以及针对 Numpy、CuPy 和 PyTorch 的特定兼容函数的关键部分是通过 array-api-compat 提供的。此软件包通过 git 子模块（位于 scipy/_lib 下）包含在 SciPy 代码库中，因此不会引入新的依赖项。

array-api-compat 提供通用实用程序函数并添加别名，例如 xp.concat（对于 numpy，映射到 np.concatenate）。这允许在 NumPy、PyTorch 和 CuPy（截至目前；预计将来会添加对 JAX 等其他库的支持）之间使用统一的 API。

当环境变量未设置，因此 SciPy 中的数组 API 标准支持被禁用时，我们仍然使用 NumPy 命名空间的“增强”版本，即 array_api_compat.numpy。这不会改变 SciPy 函数的行为，它实际上是现有的 numpy 命名空间，添加了一些别名，并为数组 API 标准支持修改/添加了一些函数。当支持启用时，根据数组的类型，xp 将返回与输入数组类型匹配的标准兼容命名空间到函数（例如，如果 cluster.vq.kmeans 的输入是 PyTorch 数组，那么 xp 是 array_api_compat.torch）。

向 SciPy 函数添加数组 API 标准支持#

尽可能地，添加到 SciPy 的新代码应该尽可能地遵循数组 API 标准（这些函数通常也是 NumPy 使用的最佳实践习惯用法）。通过遵循标准，有效地添加对数组 API 标准的支持通常很简单，理想情况下我们不需要维护任何自定义。

提供了三个辅助函数

  • array_namespace：根据输入数组返回命名空间并进行一些输入验证（例如，拒绝使用掩码数组，请参见 RFC）。

  • _asarray：asarray 的直接替换，具有附加参数 check_finite 和 order。如上所述，尝试限制非标准功能的使用。最终，我们希望将我们的需求上游到兼容性库。传递 xp=xp 可以避免内部重复调用 array_namespace。

  • copy：_asarray(x, copy=True) 的别名。copy 参数仅在 NumPy 2.0 中引入到 np.asarray，因此需要使用辅助函数来支持 <2.0。传递 xp=xp 可以避免内部重复调用 array_namespace。

要为在 .py 文件中定义的 SciPy 函数添加支持，您需要更改以下内容：

  1. 输入数组验证，

  2. 使用 xp 而不是 np 函数，

  3. 在调用编译代码时，将数组转换为 NumPy 数组，并在之后将其转换回输入数组类型。

输入数组验证使用以下模式

xp = array_namespace(arr) # where arr is the input array
# alternatively, if there are multiple array inputs, include them all:
xp = array_namespace(arr1, arr2)

# uses of non-standard parameters of np.asarray can be replaced with _asarray
arr = _asarray(arr, order='C', dtype=xp.float64, xp=xp)

请注意，如果一个输入是非 NumPy 数组类型，则所有类似数组的输入都必须是该类型；尝试将非 NumPy 数组与列表、Python 标量或其他任意 Python 对象混合使用会导致异常。对于 NumPy 数组，出于向后兼容性的原因，这些类型将继续被接受。

如果函数只调用一次编译代码，请使用以下模式

x = np.asarray(x)  # convert to numpy right before compiled call(s)
y = _call_compiled_code(x)
y = xp.asarray(y)  # convert back to original array type

如果有多次调用编译代码，请确保只进行一次转换，以避免过多的开销。

以下是一个假设的公共 SciPy 函数 toto 的示例

def toto(a, b):
    a = np.asarray(a)
    b = np.asarray(b, copy=True)

    c = np.sum(a) - np.prod(b)

    # this is some C or Cython call
    d = cdist(c)

    return d

您可以像这样转换它

def toto(a, b):
    xp = array_namespace(a, b)
    a = xp.asarray(a)
    b = copy(b, xp=xp)  # our custom helper is needed for copy

    c = xp.sum(a) - xp.prod(b)

    # this is some C or Cython call
    c = np.asarray(c)
    d = cdist(c)
    d = xp.asarray(d)

    return d

通过编译代码需要返回到 NumPy 数组，因为 SciPy 的扩展模块只与 NumPy 数组（或 Cython 中的内存视图）一起使用，而不与其他数组类型一起使用。对于 CPU 上的数组，转换应该是零拷贝的，而在 GPU 和其他设备上，转换尝试将引发异常。这样做的原因是，在设备之间进行静默数据传输被认为是不好的做法，因为它很可能是一个难以检测的大性能瓶颈。

添加测试#

以下 pytest 标记可用

  • array_api_compatible -> xp：使用参数化在多个数组后端上运行测试。

  • skip_if_array_api(*backends, reasons=None, np_only=False, cpu_only=False)：跳过某些后端和/或设备。 np_only 跳过除默认 NumPy 后端以外的所有后端的测试。 @pytest.mark.usefixtures("skip_if_array_api") 必须与该标记一起使用，才能使跳过生效。

以下是一个使用这些标记的示例

from scipy.conftest import array_api_compatible
...
@pytest.mark.skip_if_array_api(np_only=True,
                               reasons=['skip reason'])
@pytest.mark.usefixtures("skip_if_array_api")
@array_api_compatible
def test_toto1(self, xp):
    a = xp.asarray([1, 2, 3])
    b = xp.asarray([0, 2, 5])
    toto(a, b)
...
@pytest.mark.skip_if_array_api('array_api_strict', 'cupy',
                               reasons=['skip reason 1',
                                        'skip reason 2',])
@pytest.mark.usefixtures("skip_if_array_api")
@array_api_compatible
def test_toto2(self, xp):
    a = xp.asarray([1, 2, 3])
    b = xp.asarray([0, 2, 5])
    toto(a, b)

当 cpu_only=True 时，向 reasons 传递自定义原因不受支持，因为 cpu_only=True 可以与传递 backends 一起使用。此外，使用 cpu_only 的原因可能是测试的函数中使用了编译代码。

当文件中的所有测试函数都更新为兼容数组 API 后，可以使用 pytestmark 告诉 pytest 将标记应用于每个测试函数，从而减少冗余。

from scipy.conftest import array_api_compatible

pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_if_array_api")]
skip_if_array_api = pytest.mark.skip_if_array_api
...
@skip_if_array_api(np_only=True, reasons=['skip reason'])
def test_toto1(self, xp):
    a = xp.asarray([1, 2, 3])
    b = xp.asarray([0, 2, 5])
    toto(a, b)

应用这些标记后，可以使用新选项 -b 或 --array-api-backend 运行 dev.py test。

python dev.py test -b numpy -b pytorch -s cluster

这会自动设置 SCIPY_ARRAY_API。要测试具有非默认设备的多个设备的库，可以设置第二个环境变量（SCIPY_DEVICE，仅在测试套件中使用）。有效值取决于正在测试的数组库，例如，对于 PyTorch（目前唯一已知可工作的支持多设备的库），有效值为 "cpu", "cuda", "mps"。因此，要使用 PyTorch MPS 后端运行测试套件，请使用：SCIPY_DEVICE=mps python dev.py test -b pytorch。

请注意，有一个 GitHub Actions 工作流运行 pytorch-cpu。

附加信息#

以下是一些额外的资源，它们推动了一些设计决策并在开发阶段提供了帮助。

  • 包含一些讨论的初始 PR。

  • 从这个 PR 开始快速入门，并从 scikit-learn 中获得了一些灵感。

  • PR 将数组 API 支持添加到 scikit-learn。

  • 一些其他相关的 scikit-learn PR：#22554 和 #25956

本页内容