在数据科学和机器学习领域，NumPy是一个广泛使用的库，它提供了强大的数组操作功能。随着技术的发展，出现了一些新的库，它们提供了与NumPy类似的API，但可以在不同的硬件上运行，例如GPU。这些库包括CuPy和PyTorch，它们允许在GPU上进行高效的数组操作。
Scikit-learn是一个流行的机器学习库，它提供了许多用于数据挖掘和数据分析的工具。为了支持这些新的数组操作库，Scikit-learn引入了对Array API的支持。这意味着，可以在Scikit-learn中使用CuPy或PyTorch等库作为后端，进行数据操作和模型训练。
目前，这种支持还处于实验阶段，需要显式启用。只有当安装了array-api-compat库时，Scikit-learn的某些估计器才能接受Array API兼容的输入数据结构，并自动将操作派发到底层命名空间，而不是依赖于NumPy。
目前，已知与Scikit-learn的估计器兼容的Array API库包括array-api-strict、CuPy和PyTorch。下面是一个使用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中支持Array API兼容输入的估计器和其他工具包括PCA、Ridge、LinearDiscriminantAnalysis等。此外，还有一些元估计器和度量工具也支持Array API输入。随着时间的推移，支持的范围预计将会扩大。
当使用Array API兼容输入调用函数或方法时，通常的做法是返回与输入数据相同数组容器类型和设备的数组值。同样，当估计器使用Array API兼容输入进行拟合时，拟合的属性将是来自与输入相同的库的数组，并存储在相同的设备上。随后的predict和transform方法期望输入与传递给fit方法的数据来自相同的数组库和设备。
需要注意的是，返回标量值的评分函数将返回Python标量（通常是float实例），而不是数组标量值。此外，为了确保正确性并防止溢出，Scikit-learn中的某些操作会自动在float64精度的浮点值上执行操作。然而，某些数组命名空间和设备的组合（例如在MPS上的PyTorch）不支持float64数据类型。在这些情况下，Scikit-learn将改用float32数据类型。这可能导致与不使用数组API调度或使用支持float64的设备相比，行为不同（通常是数值不稳定的结果）。
要在Scikit-learn中启用对Array API的支持，需要在测试环境中安装array-api-compat库。要运行完整的检查集，需要安装PyTorch和CuPy，并拥有GPU。无法执行或缺少依赖项的检查将自动跳过。因此，运行测试时使用-v标志很重要，以便查看哪些检查被跳过。
在macOS上，PyTorch可以使用Metal Performance Shaders (MPS)来访问硬件加速器（例如M1或M2芯片内部的GPU组件）。然而，截至目前，PyTorch对MPS设备的支持还不完整。有关更多详细信息，请参见以下GitHub问题：。要在PyTorch中启用MPS支持，请在运行测试之前设置环境变量PYTORCH_ENABLE_MPS_FALLBACK=1：
PYTORCH_ENABLE_MPS_FALLBACK=1 pytest -k "array_api" -v