GridSearchCV & RandomizedSearchCV support for dpctl.tensor.usm_ndarray

[idkSeth]

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?

[napetrov]

@david-cortes-intel could add more, but here is comment from offline discussion

"But even if those two get implemented in sklearnex, the issue would still remain for other functions that perform slicing (e.g. external scikit-learn-compatible packages, plus other sklearn functions like cross_val_score).

I think the only good way of solving the issue would be to implement slicing in dpctl. Or alternatively, to use dpnp."

This is potentially something that we would consider adding, though no clear timelines yet

[david-cortes-intel]

While it's currently not possible to do this efficiently, it's still possible to run cross-validators on GPU estimators under a config context. It involves a lot more data movement, but it does achieve the goal of making all the calls to .fit() run on GPU:

import numpy as np
from sklearnex import config_context
from sklearnex.linear_model import Ridge
from sklearn.model_selection import GridSearchCV, KFold
from sklearn.datasets import make_regression

X, y = make_regression(random_state=123)
with config_context(target_offload="gpu"):
    model = GridSearchCV(
        estimator=Ridge(),
        param_grid=[{"alpha": [1e-1, 1, 1e1]}],
        cv=KFold(
            n_splits=3,
            shuffle=True,
            random_state=123,
        ),
        refit=True,
    ).fit(X, y)

# now predict on a GPU array to confirm
import dpnp
print(model.predict(dpnp.array(X, device="gpu")).device)

[syakov-intel]

IntelPython/dpctl#2053 is related to that issue