Computation backend discussion

implementation_snippet.py

        nanmin = self.get_computation_routine("nanmin")
        nanmax = self.get_computation_routine("nanmax")

        # self.get_computation_routine("nanmin_max")
        # cupy.array_api.min
        data_min = nanmin(X, axis=0)
        data_max = nanmax(X, axis=0)

        if first_pass:
            self.n_samples_seen_ = X.shape[0]
        else:
            np_min = self.get_computation_routine("min")
            np_max = self.get_computation_routine("max")
            data_min = np_min(self.data_min_, data_min)
            data_max = np_max(self.data_min_, data_min)
            self.n_samples_seen_ += X.shape[0]

        data_range = data_max - data_min
        self.scale_ = (feature_range[1] - feature_range[0]) / _handle_zeros_in_scale(
            data_range, copy=True
        )
        self.min_ = feature_range[0] - data_min * self.scale_
        self.data_min_ = data_min
        self.data_max_ = data_max
        self.data_range_ = data_range
        return self

ux_backend.py

from sklearn import set_config
from sklearn.preprocessing import MinMaxScaler
import cupy as cp
import torch
import numpy as np

cp.__sklearn_dispatchables__ = {
    "nanmin": cp.nanmin,
    "nanmax": cp.nanmax,
    "max": cp.max,
    "min": cp.nan,
}
torch.__sklearn_dispatchables__ = {
    "nanmin": torch.nanmin,
    "nanmax": torch.nanmax,
    "max": torch.max,
    "min": torch.nan,
}
set_config(backend=["cupy", "torch"])
set_config(
    backend=[
        (
            "torch",
            {
                "nanmin": torch.nanmin,
                "nanmax": torch.nanmax,
                "max": torch.max,
                "min": torch.nan,
                "clip": torch.clip,
                "finfo": torch.finfo,
                "ndarray": torch.ndarray,
                "isscalar": torch.isscalar,
            },
        )
    ]
)
set_config(global_backend="cupy")


mms = MinMaxScaler()
mms.set_computation_routine("cupy")
mms.set_computation_routine(
    {"nanmin": cp.nanmin, "nanmax": cp.nanmax, "max": cp.max, "min": cp.nan}
)

X_np = np.asarray([[1, 2, 3], [4, 5, 6]])
X_gpu = cp.asarray(X_np)
mms.fit_transform(X_gpu)

# # mms.scale_ is GPU

# mms.set_computation_routine(
#     {
#         "nanmin_max": "faster_gpu_library",
#         "min": "faster_cpu_library",
#         "max": "faster_cpu_library",
#     }
# )


pipe = make_pipeline(
    PCA().set_computation_routine("library_a"),
    MyFancyFeatureSelector().set_computation_routine("library_b"),
)


# How?
with computation_engine("library_a"):
    minmax_scale(X_gpu)


with computation_engine("torch"):
    minmax_scale(X_gpu)