12.1. 数组 API 支持（实验性）#

数组 API 规范定义了所有具有 NumPy 类似 API 的数组操作库的标准 API。Scikit-learn 提供了 array-api-compat 和 array-api-extra 的固定副本。

scikit-learn 对数组 API 标准的支持需要设置环境变量 SCIPY_ARRAY_API 为 1，并且需要在导入 scipy 和 scikit-learn 之前设置。

export SCIPY_ARRAY_API=1

请注意，此环境变量旨在临时使用。有关更多详细信息，请参阅 SciPy 的 数组 API 文档。

scikit-learn 中一些主要依赖 NumPy（而不是 Cython）来实现其 fit、predict 或 transform 方法的算法逻辑的估计器可以配置为接受任何兼容数组 API 的输入数据结构，并自动将操作调度到底层命名空间，而不是依赖 NumPy。

在此阶段，此支持被视为实验性功能，必须通过 array_api_dispatch 配置显式启用。详细信息见下文。

注意

目前，只有 array-api-strict、cupy 和 PyTorch 被确认与 scikit-learn 的估计器兼容。

以下视频概述了该标准的设计原则以及它如何促进数组库之间的互操作性

12.1.1. 启用数组 API 支持#

需要将配置 array_api_dispatch=True 设置为 True 以启用数组 API 支持。我们建议全局设置此配置，以确保行为一致并防止意外混合数组命名空间。请注意，在下面的示例中，我们使用上下文管理器（config_context）来避免在每个代码片段末尾将其重置为 False，以免影响文档的其余部分。

Scikit-learn 接受所有 metrics 和一些估计器的 array-like 输入。当 array_api_dispatch=False 时，这些输入使用 numpy.asarray（或 numpy.array）转换为 NumPy 数组。虽然这可以成功转换一些数组 API 输入（例如 JAX 数组），但我们通常建议在使用数组 API 输入时设置 array_api_dispatch=True。这是因为 NumPy 转换通常会失败，例如分配在 GPU 上的 torch tensor。

12.1.2. 使用示例#

下面的示例代码片段演示了如何使用 CuPy 在 GPU 上运行 LinearDiscriminantAnalysis

>>> from sklearn.datasets import make_classification
>>> from sklearn import config_context
>>> from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
>>> import cupy

>>> X_np, y_np = make_classification(random_state=0)
>>> X_cu = cupy.asarray(X_np)
>>> y_cu = cupy.asarray(y_np)
>>> X_cu.device
<CUDA Device 0>

>>> with config_context(array_api_dispatch=True):
...     lda = LinearDiscriminantAnalysis()
...     X_trans = lda.fit_transform(X_cu, y_cu)
>>> X_trans.device
<CUDA Device 0>

模型训练完成后，作为数组的拟合属性也将来自与训练数据相同的数组 API 命名空间。例如，如果使用 CuPy 的数组 API 命名空间进行训练，则拟合属性将位于 GPU 上。我们提供了一个实验性的 _estimator_with_converted_arrays 工具，用于将估计器属性从数组 API 传输到 ndarray

>>> from sklearn.utils._array_api import _estimator_with_converted_arrays
>>> cupy_to_ndarray = lambda array : array.get()
>>> lda_np = _estimator_with_converted_arrays(lda, cupy_to_ndarray)
>>> X_trans = lda_np.transform(X_np)
>>> type(X_trans)
<class 'numpy.ndarray'>

12.1.2.1. PyTorch 支持#

PyTorch Tensors 也可以直接传入

>>> import torch
>>> X_torch = torch.asarray(X_np, device="cuda", dtype=torch.float32)
>>> y_torch = torch.asarray(y_np, device="cuda", dtype=torch.float32)

>>> with config_context(array_api_dispatch=True):
...     lda = LinearDiscriminantAnalysis()
...     X_trans = lda.fit_transform(X_torch, y_torch)
>>> type(X_trans)
<class 'torch.Tensor'>
>>> X_trans.device.type
'cuda'

12.1.3. 支持 Array API 兼容输入的组件#

scikit-learn 中支持数组 API 兼容输入的估计器和其他工具。

12.1.3.1. 估计器#

12.1.3.2. 元估计器#

接受数组 API 输入的元估计器，前提是基础估计器也支持。

12.1.3.3. 评估指标#

12.1.3.4. 工具#

覆盖范围预计会随着时间而增加。请关注 GitHub 上的专用 meta-issue 以跟踪进度。

12.1.4. 输入和输出数组类型处理#

估计器和评分函数能够接受来自不同数组库和/或设备的输入数组。当传入一组混合的输入数组时，scikit-learn 会根据需要转换数组以使它们全部保持一致。

对于估计器，规则是**“一切跟随 X”** - 混合数组输入被转换，以使它们都与 X 的数组库和设备匹配。对于评分函数，规则是**“一切跟随 y_pred”** - 混合数组输入被转换，以使它们都与 y_pred 的数组库和设备匹配。

当使用兼容数组 API 的输入调用函数或方法时，约定是返回与输入数据来自相同数组库且位于相同设备上的数组。

12.1.4.1. 估计器#

当估计器使用兼容数组 API 的 X 进行拟合时，所有其他数组输入，包括构造函数参数（例如 y、sample_weight），如果它们不匹配，将被转换以匹配 X 的数组库和设备。这种行为使得可以在管道中的任何点从 CPU 处理切换到 GPU 处理。

这允许估计器接受混合输入类型，从而使 X 可以在管道中移动到不同的设备，而无需显式移动 y。请注意，scikit-learn 管道不允许转换 y（以避免数据泄露）。

以一个管道为例，其中 X 和 y 都从 CPU 开始，并经历以下三个步骤

X 最初包含分类字符串数据（因此需要在 CPU 上），它在 TargetEncoder 中被目标编码为数值。然后 X 被显式移动到 GPU，以提高 Ridge 的性能。y 不能被管道转换（回想一下 scikit-learn 管道不允许转换 y），但由于 Ridge 能够接受混合输入类型，这不是问题，管道能够运行。

使用兼容数组 API 的 X 拟合的估计器的拟合属性将是来自相同库并存储在相同设备上的数组。随后的 predict 和 transform 方法需要来自与传递给 fit 方法的数据相同的数组库和设备上的输入。

12.1.4.2. 评分函数#

当兼容数组 API 的 y_pred 传递给评分函数时，所有其他数组输入（例如 y_true、sample_weight），如果它们不匹配，将被转换以匹配 y_pred 的数组库和设备。这允许评分函数接受混合输入类型，从而使它们可以在元估计器（或接受估计器的函数）中使用，其中管道在设备之间移动输入数组（例如，CPU 到 GPU）。

例如，为了能够在 cross_validate 或 GridSearchCV 等中使用上述管道，内部调用的评分函数需要能够接受混合输入类型。

评分函数的输出类型取决于输出值的数量。当评分函数返回标量值时，它将返回 Python 标量（通常是 float 实例）而不是数组标量值。对于支持多类或多输出的评分函数，当需要输出多个值时，将返回与 y_pred 来自相同数组库和设备上的数组。

12.1.5. 常见估计器检查#

将 array_api_support 标签添加到估计器的标签集中，以指示它支持数组 API。这将启用专用检查作为常见测试的一部分，以验证估计器在使用普通 NumPy 和数组 API 输入时产生相同的结果。

要运行这些检查，您需要在测试环境中安装 array-api-strict。这使您无需 GPU 即可运行检查。要运行全套检查，您还需要安装 PyTorch、CuPy 并拥有 GPU。无法执行或缺少依赖项的检查将自动跳过。因此，运行测试时使用 -v 标志很重要，以查看跳过了哪些检查

pip install array-api-strict  # and other libraries as needed
pytest -k "array_api" -v

针对 array-api-strict 运行 scikit-learn 测试应该有助于发现大多数与通过使用模拟的非 CPU 设备处理多个设备输入相关的代码问题。这使得数组 API 相关代码的快速迭代开发和调试成为可能。

但是，为了确保完全处理分配在实际 GPU 设备上的 PyTorch 或 CuPy 输入，有必要针对这些库和硬件运行测试。这可以通过使用 Google Colab 或利用拉取请求上的 CI 基础设施（出于成本原因由维护人员手动触发）来实现。

12.1.5.1. MPS 设备支持说明#

在 macOS 上，PyTorch 可以使用 Metal Performance Shaders (MPS) 来访问硬件加速器（例如 M1 或 M2 芯片的内部 GPU 组件）。然而，在撰写本文时，PyTorch 对 MPS 设备的支持并不完整。有关更多详细信息，请参阅以下 github issue

要在 PyTorch 中启用 MPS 支持，请在运行测试之前设置环境变量 PYTORCH_ENABLE_MPS_FALLBACK=1

PYTORCH_ENABLE_MPS_FALLBACK=1 pytest -k "array_api" -v

在撰写本文时，所有 scikit-learn 测试都应该通过，但是，计算速度不一定比使用 CPU 设备快。

12.1.5.2. float64 设备支持说明#

scikit-learn 中的某些操作会自动对浮点值执行 float64 精度操作，以防止溢出并确保正确性（例如，metrics.pairwise.euclidean_distances、preprocessing.StandardScaler）。然而，某些数组命名空间和设备的组合，例如 PyTorch on MPS（参见 MPS 设备支持说明），不支持 float64 数据类型。在这些情况下，scikit-learn 将转而使用 float32 数据类型。这可能导致与不使用数组 API 调度或使用支持 float64 的设备相比不同的行为（通常是数值不稳定的结果）。