在机器学习领域，数据处理和模型训练往往需要高效的数值计算库。NumPy作为Python中广泛使用的数值计算库，提供了强大的数组操作功能。然而，随着硬件技术的发展，GPU成为了加速计算的重要工具。为了充分利用GPU的计算能力，scikit-learn库开始支持与Array API兼容的输入数据结构，允许用户在GPU上进行模型的训练和预测。
目前，scikit-learn的这一支持还处于实验阶段，需要用户显式启用。通过安装array-api-compat库，用户可以让scikit-learn的某些估计器接受Array API兼容的输入数据结构，并自动将操作调度到底层命名空间，而不是依赖于NumPy。这意味着，如果用户使用CuPy或PyTorch等库作为底层数组操作库，那么scikit-learn的估计器将能够在GPU上执行计算任务。
例如，使用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>
在模型训练完成后，拟合的属性（如数组）将与训练数据来自同一Array API命名空间。例如，如果使用CuPy的Array API命名空间进行训练，那么拟合的属性将位于GPU上。scikit-learn提供了一个实验性的_estimator_with_converted_arrays工具，可以将估计器属性从Array 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'>
对于PyTorch的支持，用户可以通过设置array_api_dispatch=True并直接传递张量来实现：
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'
目前，scikit-learn支持的估计器和其他工具包括PCA、Ridge、LinearDiscriminantAnalysis、KernelCenterer、MaxAbsScaler、MinMaxScaler、Normalizer等。此外，还支持一些度量工具，如accuracy_score、cosine_similarity、r2_score、zero_one_loss等。随着时间的推移，预计支持的范围将不断扩大。用户可以关注GitHub上的专门问题来跟踪进展。
当使用Array API兼容的输入调用函数或方法时，约定是返回与输入数据相同数组容器类型和设备的数组值。同样，当估计器使用Array API兼容的输入进行拟合时，拟合的属性将是来自与输入相同的库的数组，并存储在相同的设备上。随后的predict和transform方法期望输入与传递给fit方法的数据来自相同的数组库和设备。需要注意的是，返回标量值的评分函数返回的是Python标量（通常是float实例），而不是数组标量值。
为了表示估计器支持Array API，可以在估计器的标签集中添加array_api_support标签。这将启用专门的检查，作为通用测试的一部分，以验证当使用普通的NumPy和Array API输入时，估计器的结果是否相同。要运行这些检查，需要在测试环境中安装array_api_compat库。要运行完整的检查集，需要安装PyTorch和CuPy，并拥有GPU。无法执行或缺少依赖项的检查将自动跳过。因此，运行测试时使用-v标志很重要，以查看哪些检查被跳过：
pip install array-api-compat # 和其他需要的库
pytest -k "array_api" -v
关于MPS设备支持的说明：在macOS上，PyTorch可以使用Metal Performance Shaders (MPS)来访问硬件加速器（例如M1或M2芯片的内部GPU组件）。然而，截至本文撰写时，PyTorch对MPS设备的支持还不完整。有关更多详细信息，请参见以下GitHub问题：pytorch/pytorch#77764。要在PyTorch中启用MPS支持，请在运行测试之前设置环境变量PYTORCH_ENABLE_MPS_FALLBACK=1：
PYTORCH_ENABLE_MPS_FALLBACK=1 pytest -k "array_api" -v