支持数组 API 标准#
注意
数组 API 标准支持仍处于实验阶段，并且隐藏在环境变量背后。目前只覆盖了公共 API 的一小部分。
本指南描述了如何使用和添加对 Python 数组 API 标准 的支持。此标准允许用户直接使用任何与数组 API 兼容的数组库和 SciPy 的部分功能。
RFC 定义了 SciPy 如何实现对该标准的支持，其主要原则是“输入数组类型等于输出数组类型”。此外，该实现还对允许的类数组输入进行了更严格的验证，例如拒绝 NumPy 矩阵和掩码数组实例，以及具有对象 dtype 的数组。
在下文中，与数组 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 在其实现中使用了编译代码，这在 GPU 上无法工作。
更严格的数组输入验证将拒绝 np.matrix 和 np.ma.MaskedArray 实例，以及具有 object dtype 的数组
>>> 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 标准支持
在 scipy.special 中为以下函数提供了支持：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, scipy.special.expit, scipy.special.entr, scipy.special.rel_entr, scipy.special.rel_entr, scipy.special.xlogy, 以及 scipy.special.chdtrc。
在 scipy.stats 模块中，为以下函数提供了支持：scipy.stats.describe、scipy.stats.moment、scipy.stats.skew、scipy.stats.kurtosis、scipy.stats.kstat、scipy.stats.kstatvar、scipy.stats.circmean、scipy.stats.circvar、scipy.stats.circstd、scipy.stats.entropy、scipy.stats.variation 、scipy.stats.sem、scipy.stats.ttest_1samp、scipy.stats.pearsonr、scipy.stats.chisquare、scipy.stats.skewtest、scipy.stats.kurtosistest、scipy.stats.normaltest、scipy.stats.jarque_bera、scipy.stats.bartlett、scipy.stats.power_divergence 和 scipy.stats.monte_carlo_test。
请查看跟踪问题以获取更新。
实现说明#
通过 array-api-compat 提供了对数组 API 标准和 Numpy、CuPy 和 PyTorch 的特定兼容性函数的关键支持。此软件包通过 git 子模块（在 scipy/_lib 下）包含在 SciPy 代码库中，因此不会引入新的依赖项。
array-api-compat 提供了通用实用程序函数并添加了别名，例如 xp.concat（对于 numpy，在 NumPy 2.0 中添加 np.concat 之前，映射到 np.concatenate）。这允许在 NumPy、PyTorch、CuPy 和 JAX（以及正在开发的 Dask 等其他库）中使用统一的 API。
当未设置环境变量，因此 SciPy 中禁用数组 API 标准支持时，我们仍然使用 NumPy 命名空间的包装版本，即 array_api_compat.numpy。这不应更改 SciPy 函数的行为，因为它实际上是现有的 numpy 命名空间，添加了许多别名，并为数组 API 标准支持修改/添加了一些函数。启用支持后，xp = array_namespace(input) 将是与输入数组类型匹配的、与标准兼容的命名空间（例如，如果 cluster.vq.kmeans 的输入是 PyTorch 张量，则 xp 将是 array_api_compat.torch）。
向 SciPy 函数添加数组 API 标准支持#
尽可能地，添加到 SciPy 的新代码应尽可能地遵循数组 API 标准（这些函数通常也是 NumPy 用法的最佳实践习惯用法）。通过遵循标准，有效地添加对数组 API 标准的支持通常很简单，并且我们理想情况下不需要维护任何自定义项。
scipy._lib._array_api 中提供了各种帮助程序函数 - 请查看该模块中的 __all__ 以获取当前帮助程序的列表，并查看其文档字符串以获取更多信息。
要向 .py 文件中定义的 SciPy 函数添加支持，您必须更改的是
输入数组验证，
使用
xp而不是np函数，当调用编译的代码时，先将数组转换为 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)
# replace np.asarray with xp.asarray
arr = xp.asarray(arr)
# 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 = xp_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 中的 memoryviews）。对于 CPU 上的数组，转换应该是零拷贝的，而在 GPU 和其他设备上，转换尝试将引发异常。原因是设备之间的数据静默传输被认为是不好的做法，因为它很可能是一个大的且难以检测的性能瓶颈。
添加测试#
以下 pytest 标记可用
array_api_compatible -> xp：使用参数化在多个数组后端上运行测试。skip_xp_backends(backend=None, reason=None, np_only=False, cpu_only=False, exceptions=None)：跳过某些后端或后端类别。必须将@pytest.mark.usefixtures("skip_xp_backends")与此标记一起使用才能应用跳过。有关如何使用此标记跳过测试的信息，请参阅scipy.conftest中的 fixture 的文档字符串。xfail_xp_backends(backend=None, reason=None, np_only=False, cpu_only=False, exceptions=None)：使某些后端或后端类别预期失败。必须将@pytest.mark.usefixtures("xfail_xp_backends")与此标记一起使用才能应用预期失败。有关如何使用此标记使测试预期失败的信息，请参阅scipy.conftest中的 fixture 的文档字符串。当启用
SCIPY_ARRAY_API时，skip_xp_invalid_arg用于跳过使用无效参数的测试。例如，scipy.stats函数的一些测试将掩码数组传递给正在测试的函数，但掩码数组与数组 API 不兼容。使用skip_xp_invalid_arg装饰器允许这些测试在未使用SCIPY_ARRAY_API时防止回归，而不会在使用SCIPY_ARRAY_API时导致失败。随着时间的推移，我们希望这些函数在接收到数组 API 无效输入时发出弃用警告，并且此装饰器将检查是否发出弃用警告而不会导致测试失败。当SCIPY_ARRAY_API=1行为成为默认行为且唯一行为时，将删除这些测试（和装饰器本身）。
scipy._lib._array_api 包含与数组无关的断言，例如 xp_assert_close，它可用于替换 numpy.testing 中的断言。
以下示例演示如何使用标记
from scipy.conftest import array_api_compatible, skip_xp_invalid_arg
from scipy._lib._array_api import xp_assert_close
...
@pytest.mark.skip_xp_backends(np_only=True, reason='skip reason')
@pytest.mark.usefixtures("skip_xp_backends")
@array_api_compatible
def test_toto1(self, xp):
a = xp.asarray([1, 2, 3])
b = xp.asarray([0, 2, 5])
xp_assert_close(toto(a, b), a)
...
@pytest.mark.skip_xp_backends('array_api_strict',
reason='skip reason 1')
@pytest.mark.skip_xp_backends('cupy',
reason='skip reason 2')
@pytest.mark.usefixtures("skip_xp_backends")
@array_api_compatible
def test_toto2(self, xp):
...
...
# Do not run when SCIPY_ARRAY_API is used
@skip_xp_invalid_arg
def test_toto_masked_array(self):
...
当 cpu_only=True 时，将自定义原因传递给 reason 不受支持，因为 cpu_only=True 可以与传递 backends 一起使用。此外，使用 cpu_only 的原因可能只是因为编译的代码在正在测试的函数中使用。
将后端名称传递给 exceptions 意味着它们不会被 cpu_only=True 跳过。当为某些（但不是全部）非 CPU 后端实现了委托，并且 CPU 代码路径需要转换为 NumPy 以进行编译代码时，这很有用。
# array-api-strict and CuPy will always be skipped, for the given reasons.
# All libraries using a non-CPU device will also be skipped, apart from
# JAX, for which delegation is implemented (hence non-CPU execution is supported).
@pytest.mark.skip_xp_backends(cpu_only, exceptions=['jax.numpy'])
@pytest.mark.skip_xp_backends('array_api_strict', reason='skip reason 1')
@pytest.mark.skip_xp_backends('cupy', reason='skip reason 2')
@pytest.mark.usefixtures("skip_xp_backends")
@array_api_compatible
def test_toto(self, xp):
...
当文件中的每个测试函数都已更新以兼容 Array API 时，可以通过告诉 pytest 使用 pytestmark 将标记应用于每个测试函数来减少冗长。
from scipy.conftest import array_api_compatible
pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
skip_xp_backends = pytest.mark.skip_xp_backends
...
@skip_xp_backends(np_only=True, reason='skip reason')
def test_toto1(self, xp):
...
应用这些标记后，可以使用新选项 -b 或 --array-api-backend 来使用 dev.py test。
python dev.py test -b numpy -b torch -s cluster
这将自动设置 SCIPY_ARRAY_API。要测试具有非默认设备的多个设备的库，可以设置第二个环境变量（SCIPY_DEVICE，仅在测试套件中使用）。有效值取决于被测的数组库，例如，对于 PyTorch，有效值是 "cpu"、"cuda"、"mps"。要使用 PyTorch MPS 后端运行测试套件，请使用： SCIPY_DEVICE=mps python dev.py test -b torch。
请注意，有一个 GitHub Actions 工作流程，它在 CPU 上使用 array-api-strict、PyTorch 和 JAX 进行测试。
其他信息#
以下是一些其他资源，这些资源推动了一些设计决策并在开发阶段提供了帮助