Calcul GPU natif via Array API

Scikit-learn 1.8 étend le support de l'Array API standard, permettant de passer directement des tableaux PyTorch ou CuPy aux algorithmes supportés. Le calcul reste sur le GPU sans copie vers NumPy, offrant des accélérations significatives sur les grands jeux de données.

Exemple avec PyTorch

python
import torch
from sklearn.linear_model import Ridge
from sklearn import config_context

# Données directement sur GPU
X = torch.randn(10000, 100, device='cuda')
y = torch.randn(10000, device='cuda')

# Activer le dispatch Array API
with config_context(array_api_dispatch=True):
    model = Ridge(alpha=1.0)
    model.fit(X, y)
    predictions = model.predict(X)

print(type(predictions))  # torch.Tensor
print(predictions.device)  # cuda:0

# Algorithmes supportés : Ridge, KMeans, PCA,
# LinearDiscriminantAnalysis, etc.

Sources