GridSearchCV & RandomizedSearchCV GPU support

[napetrov]

Discussed in #2350

^{Originally posted by idkSeth March 8, 2025}
In order to offload to GPU, input data needs to be passed as dpctl.tensor.usm_ndarray as per the guide.

GridSearchCV and RandomizedSearchCV is not patched and when trying to use them with the input data as dpctl.tensor.usm_ndarray, an indexing error is raised.

Can searching be patched to allow for hyper-parameter searching on GPU/CPU with usm_ndarray?

[david-cortes-intel]

This is now solved on the dpctl side: IntelPython/dpctl#2128

Should be possible to use the classes from sklearn with sklearnex estimators and dpctl/dpnp arrays after version 0.21 of dpctl.

[david-cortes-intel]

This is now possible to do with array API classes.

Example with DPNP (>= 0.19):

import os
os.environ["SKLEARNEX_VERBOSE"] = "INFO"
os.environ["SCIPY_ARRAY_API"] = "1"
import numpy as np
import dpnp
from sklearnex import config_context
from sklearnex.linear_model import Ridge
from sklearn.datasets import make_regression
from sklearn.model_selection import GridSearchCV, KFold

X_cpu, y_cpu = make_regression(random_state=123)
X_gpu = dpnp.array(X_cpu, dtype=np.float32, device="gpu")
y_gpu = dpnp.array(y_cpu, dtype=np.float32, device="gpu")

cv = GridSearchCV(
    estimator=Ridge(),
    param_grid={"alpha": [1e-1, 1, 1e1]},
    n_jobs=1,
    refit=True,
    cv=KFold(
        n_splits=5,
        shuffle=True,
        random_state=123,
    ),
)
with config_context(array_api_dispatch=True):
    cv.fit(X_gpu, y_gpu)

Example with PyTorch:

import os
os.environ["SKLEARNEX_VERBOSE"] = "INFO"
os.environ["SCIPY_ARRAY_API"] = "1"
import torch
from sklearnex import config_context
from sklearnex.linear_model import Ridge
from sklearn.datasets import make_regression
from sklearn.model_selection import GridSearchCV, KFold

X_cpu, y_cpu = make_regression(random_state=123)
X_gpu = torch.tensor(X_cpu, dtype=torch.float32, device="xpu")
y_gpu = torch.tensor(y_cpu, dtype=torch.float32, device="xpu")

cv = GridSearchCV(
    estimator=Ridge(),
    param_grid={"alpha": [1e-1, 1, 1e1]},
    n_jobs=1,
    refit=True,
    cv=KFold(
        n_splits=5,
        shuffle=True,
        random_state=123,
    ),
)
with config_context(array_api_dispatch=True):
    cv.fit(X_gpu, y_gpu)