Native GPU Computing via Array API

Scikit-learn 1.8 extends Array API standard support, allowing you to pass PyTorch or CuPy arrays directly to supported algorithms. Computation stays on GPU without copying to NumPy, offering significant speedups on large datasets.

Example with PyTorch

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

# Data directly on GPU
X = torch.randn(10000, 100, device='cuda')
y = torch.randn(10000, device='cuda')

# Enable Array API dispatch
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

# Supported algorithms: Ridge, KMeans, PCA,
# LinearDiscriminantAnalysis, etc.

Sources