数组 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 标准支持
在 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
。
实现说明#
对数组 API 标准的支持以及对 Numpy、CuPy 和 PyTorch 的特定兼容函数的关键部分是通过 array-api-compat 提供的。此包通过 git 子模块(位于 scipy/_lib
下)包含在 SciPy 代码库中,因此不会引入新的依赖项。
array-api-compat
提供通用实用程序函数并添加别名,例如 xp.concat
(对于 numpy,它映射到 np.concatenate
)。这允许在 NumPy、PyTorch、CuPy 和 JAX(以及 Dask 等其他库,将在未来推出)之间使用统一的 API。
当环境变量未设置,因此 SciPy 中的数组 API 标准支持被禁用时,我们仍然使用 NumPy 命名空间的“增强”版本,即 array_api_compat.numpy
。这不会改变 SciPy 函数的行为,它实际上是现有的 numpy
命名空间,添加了一些别名,并为数组 API 标准支持添加了一些函数。当支持启用时,根据数组类型,xp
将返回与输入数组类型匹配的标准兼容命名空间到函数(例如,如果 cluster.vq.kmeans 的输入是 PyTorch 数组,则 xp
是 array_api_compat.torch
)。
将数组 API 标准支持添加到 SciPy 函数#
尽可能地,添加到 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 函数中,您需要更改的内容是
输入数组验证,
使用
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)
# 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_xp_backends(*backends, reasons=None, np_only=False, cpu_only=False)
:跳过某些后端和/或设备。np_only
跳过除默认 NumPy 后端之外所有后端的测试。@pytest.mark.usefixtures("skip_xp_backends")
必须与该标记一起使用,才能应用跳过操作。skip_xp_invalid_arg
用于跳过在使用SCIPY_ARRAY_API
时使用无效参数的测试。例如,scipy.stats
函数的一些测试将掩码数组传递给被测试的函数,但掩码数组与数组 API 不兼容。使用skip_xp_invalid_arg
装饰器允许这些测试在未使用SCIPY_ARRAY_API
时防止回归,而不会在使用SCIPY_ARRAY_API
时导致失败。随着时间的推移,我们希望这些函数在接收数组 API 无效输入时发出弃用警告,并且此装饰器将检查是否发出弃用警告,而不会导致测试失败。当SCIPY_ARRAY_API=1
行为成为默认行为和唯一行为时,这些测试(以及装饰器本身)将被删除。
以下是用标记的示例
from scipy.conftest import array_api_compatible, skip_xp_invalid_arg
...
@pytest.mark.skip_xp_backends(np_only=True,
reasons=['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])
toto(a, b)
...
@pytest.mark.skip_xp_backends('array_api_strict', 'cupy',
reasons=['skip reason 1',
'skip reason 2',])
@pytest.mark.usefixtures("skip_xp_backends")
@array_api_compatible
def test_toto2(self, xp):
a = xp.asarray([1, 2, 3])
b = xp.asarray([0, 2, 5])
toto(a, b)
...
# Do not run when SCIPY_ARRAY_API is used
@skip_xp_invalid_arg
def test_toto_masked_array(self):
a = np.ma.asarray([1, 2, 3])
b = np.ma.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_xp_backends")]
skip_xp_backends = pytest.mark.skip_xp_backends
...
@skip_xp_backends(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
。
其他信息#
以下是一些其他资源,它们激发了某些设计决策,并在开发阶段提供了帮助