ENH: DBSCAN via Array API

onedal/cluster/dbscan.py

@@ -14,14 +14,14 @@
 # limitations under the License.
 # ===============================================================================
 
-import numpy as np
+from sklearn.utils import check_array
 
-from daal4py.sklearn._utils import get_dtype, make2d
+from onedal.datatypes._data_conversion import get_dtype, make2d
 
 from ..common._base import BaseEstimator
 from ..common._mixin import ClusterMixin
 from ..datatypes import _convert_to_supported, from_table, to_table
-from ..utils import _check_array
+from ..utils._array_api import get_namespace
 
 
 class BaseDBSCAN(BaseEstimator, ClusterMixin):
@@ -46,38 +46,54 @@ def __init__(
         self.p = p
         self.n_jobs = n_jobs
 
-    def _get_onedal_params(self, dtype=np.float32):
+    def _get_onedal_params(self, xp, dtype):
         return {
-            "fptype": "float" if dtype == np.float32 else "double",
+            "fptype": "float" if dtype == xp.float32 else "double",
             "method": "by_default",
             "min_observations": int(self.min_samples),
             "epsilon": float(self.eps),
             "mem_save_mode": False,
             "result_options": "core_observation_indices|responses",
         }
 
-    def _fit(self, X, y, sample_weight, module, queue):
+    def _fit(self, X, xp, is_array_api_compliant, y, sample_weight, queue):
         policy = self._get_policy(queue, X)
-        X = _check_array(X, accept_sparse="csr", dtype=[np.float64, np.float32])
+        X = check_array(X, accept_sparse="csr", dtype=[xp.float64, xp.float32])
         sample_weight = make2d(sample_weight) if sample_weight is not None else None
         X = make2d(X)
+        if xp:
+            X_device = X.device
 
-        types = [np.float32, np.float64]
+        # TODO:
+        # revice once again the flow.
+        types = [xp.float32, xp.float64]
         if get_dtype(X) not in types:
-            X = X.astype(np.float64)
-        X = _convert_to_supported(policy, X)
+            X = X.astype(xp.float64)
+        X = _convert_to_supported(policy, X, xp)
+        sample_weight = (
+            _convert_to_supported(policy, sample_weight, xp)
+            if sample_weight is not None
+            else None
+        )
         dtype = get_dtype(X)
-        params = self._get_onedal_params(dtype)
-        result = module.compute(policy, params, to_table(X), to_table(sample_weight))
+        params = self._get_onedal_params(xp, dtype)
+        result = self._get_backend("dbscan", "clustering", None).compute(
+            policy, params, to_table(X, xp), to_table(sample_weight, xp)
+        )
 
-        self.labels_ = from_table(result.responses).ravel()
+        self.labels_ = from_table(
+            result.responses, xp=xp, queue=queue, array_api_device=X_device
+        ).reshape(-1)
         if result.core_observation_indices is not None:
             self.core_sample_indices_ = from_table(
-                result.core_observation_indices
-            ).ravel()
+                result.core_observation_indices,
+                xp=xp,
+                queue=queue,
+                array_api_device=X_device,
+            ).reshape(-1)
         else:
-            self.core_sample_indices_ = np.array([], dtype=np.intc)
-        self.components_ = np.take(X, self.core_sample_indices_, axis=0)
+            self.core_sample_indices_ = xp.array([], dtype=xp.int32)
+        self.components_ = xp.take(X, self.core_sample_indices_, axis=0)
         self.n_features_in_ = X.shape[1]
         return self
 
@@ -105,6 +121,5 @@ def __init__(
         self.n_jobs = n_jobs
 
     def fit(self, X, y=None, sample_weight=None, queue=None):
-        return super()._fit(
-            X, y, sample_weight, self._get_backend("dbscan", "clustering", None), queue
-        )
+        xp, is_array_api_compliant = get_namespace(X)
+        return super()._fit(X, xp, is_array_api_compliant, y, sample_weight, queue)

onedal/datatypes/_data_conversion.py

@@ -18,48 +18,76 @@
 
 import numpy as np
 
-from daal4py.sklearn._utils import make2d
+from daal4py.sklearn._utils import get_dtype
+from daal4py.sklearn._utils import make2d as d4p_make2d
 from onedal import _backend, _is_dpc_backend
 
 from ..utils import _is_csr
 
-try:
-    import dpctl
-    import dpctl.tensor as dpt
+# from ..utils._array_api import get_namespace
+from ..utils._array_api import _is_numpy_namespace
 
-    dpctl_available = dpctl.__version__ >= "0.14"
-except ImportError:
-    dpctl_available = False
 
+# TODO:
+# move to proper module.
+# TODO
+# def make2d(arg, xp=None, is_array_api_compliant=None):
+def make2d(arg, xp=None):
+    if xp and not _is_numpy_namespace(xp) and arg.ndim == 1:
+        return xp.reshape(arg, (arg.size, 1)) if arg.ndim == 1 else arg
+    # TODO:
+    # reimpl via is_array_api_compliant usage.
+    return d4p_make2d(arg)
 
-def _apply_and_pass(func, *args):
+
+# TODO:
+# remove such kind of func calls
+def _apply_and_pass(func, *args, **kwargs):
+    # TODO:
+    # refactor.
     if len(args) == 1:
         return func(args[0])
-    return tuple(map(func, args))
-
+    return tuple(map(func, args, kwargs))
 
-def from_table(*args):
-    return _apply_and_pass(_backend.from_table, *args)
 
+def convert_one_from_table(arg, xp=None, queue=None, array_api_device=None):
+    # TODO:
+    # use `array_api_device`.
+    result = _backend.from_table(arg)
+    if xp:
+        if xp.__name__ in {"dpctl", "dpctl.tensor"}:
+            result = xp.asarray(arg, sycl_queue=queue) if queue else xp.asarray(arg)
+        elif not _is_numpy_namespace(xp):
+            results = xp.asarray(result)
+    return result
 
-def convert_one_to_table(arg):
-    if dpctl_available:
-        if isinstance(arg, dpt.usm_ndarray):
-            return _backend.dpctl_to_table(arg)
 
+def convert_one_to_table(arg, xp=None):
     if not _is_csr(arg):
+        if xp and not _is_numpy_namespace(xp):
+            arg = np.asarray(arg)
+        # TODO:
+        # Check. Probably should be removed from here
+        # Not really realted with converting to table.
         arg = make2d(arg)
     return _backend.to_table(arg)
 
 
-def to_table(*args):
-    return _apply_and_pass(convert_one_to_table, *args)
+def from_table(*args, xp=None, queue=None, array_api_device=None):
+    return _apply_and_pass(convert_one_from_table, *args)
+
+
+def to_table(*args, xp=None):
+    return _apply_and_pass(convert_one_to_table, *args, xp=xp)
 
 
 if _is_dpc_backend:
     from ..common._policy import _HostInteropPolicy
 
-    def _convert_to_supported(policy, *data):
+    def _convert_to_supported(policy, *data, xp=None):
+        if xp is None:
+            xp = np
+
         def func(x):
             return x
 
@@ -71,13 +99,13 @@ def func(x):
         device = policy._queue.sycl_device
 
         def convert_or_pass(x):
-            if (x is not None) and (x.dtype == np.float64):
+            if (x is not None) and (x.dtype == xp.float64):
                 warnings.warn(
                     "Data will be converted into float32 from "
                     "float64 because device does not support it",
                     RuntimeWarning,
                 )
-                return x.astype(np.float32)
+                return x.astype(xp.float32)
             else:
                 return x
 
@@ -88,7 +116,7 @@ def convert_or_pass(x):
 
 else:
 
-    def _convert_to_supported(policy, *data):
+    def _convert_to_supported(policy, *data, xp=None):
         def func(x):
             return x
 

onedal/utils/_array_api.py

@@ -18,6 +18,9 @@
 
 from collections.abc import Iterable
 
+import numpy as np
+
+
 try:
     from dpctl.tensor import usm_ndarray
 
@@ -89,3 +92,25 @@ def _get_sycl_namespace(*arrays):
             raise ValueError(f"SYCL type not recognized: {sycl_type}")
 
     return sycl_type, None, False
+
+
+def get_namespace(*arrays):
+    """Get namespace of arrays.
+    TBD.
+    Parameters
+    ----------
+    *arrays : array objects
+        Array objects.
+    Returns
+    -------
+    namespace : module
+        Namespace shared by array objects.
+    is_array_api : bool
+        True of the arrays are containers that implement the Array API spec.
+    """
+    sycl_type, xp, is_array_api_compliant = _get_sycl_namespace(*arrays)
+
+    if sycl_type:
+        return xp, is_array_api_compliant
+    else:
+        return np, True

sklearnex/_device_offload.py

@@ -61,6 +61,40 @@ def _get_backend(obj, queue, method_name, *data):
     raise RuntimeError("Device support is not implemented")
 
 
+# def dispatch_with_array_api(obj, method_name, branches, xp, is_array_api_compliant, *args, **kwargs):
+def dispatch_with_array_api(
+    obj, method_name, branches, xp, is_array_api_compliant, *args, **kwargs
+):
+    q = _get_global_queue()
+    # if "array_api_support_sklearnex" in obj._get_tags() and obj._get_tags()["array_api_support_sklearnex"]:
+
+    backend, q, patching_status = _get_backend(obj, q, method_name, *args)
+
+    if backend == "onedal":
+        patching_status.write_log(queue=q)
+        return branches[backend](obj, *args, **kwargs, queue=q)
+    if backend == "sklearn":
+        if (
+            "array_api_dispatch" in get_config()
+            and get_config()["array_api_dispatch"]
+            and "array_api_support" in obj._get_tags()
+            and obj._get_tags()["array_api_support"]
+        ):
+            # If `array_api_dispatch` enabled and array api is supported for the stock scikit-learn,
+            # then raw inputs are used for the fallback.
+            patching_status.write_log()
+            return branches[backend](obj, *args, **kwargs)
+        else:
+            patching_status.write_log()
+            _, hostargs = _transfer_to_host(q, *args)
+            _, hostvalues = _transfer_to_host(q, *kwargs.values())
+            hostkwargs = dict(zip(kwargs.keys(), hostvalues))
+            return branches[backend](obj, *hostargs, **hostkwargs)
+    raise RuntimeError(
+        f"Undefined backend {backend} in " f"{obj.__class__.__name__}.{method_name}"
+    )
+
+
 def dispatch(obj, method_name, branches, *args, **kwargs):
     q = _get_global_queue()
     q, hostargs = _transfer_to_host(q, *args)

sklearnex/cluster/dbscan.py

@@ -25,8 +25,9 @@
 from daal4py.sklearn._utils import sklearn_check_version
 from onedal.cluster import DBSCAN as onedal_DBSCAN
 
-from .._device_offload import dispatch
+from .._device_offload import dispatch, dispatch_with_array_api
 from .._utils import PatchingConditionsChain
+from ..utils._array_api import get_namespace
 
 if sklearn_check_version("1.1") and not sklearn_check_version("1.2"):
     from sklearn.utils import check_scalar
@@ -91,6 +92,7 @@ def __init__(
     def _onedal_fit(self, X, y, sample_weight=None, queue=None):
         if sklearn_check_version("1.0"):
             X = validate_data(self, X, force_all_finite=False)
+        xp, is_array_api_compliant = get_namespace(X)
 
         onedal_params = {
             "eps": self.eps,
@@ -104,7 +106,9 @@ def _onedal_fit(self, X, y, sample_weight=None, queue=None):
         }
         self._onedal_estimator = self._onedal_dbscan(**onedal_params)
 
-        self._onedal_estimator.fit(X, y=y, sample_weight=sample_weight, queue=queue)
+        self._onedal_estimator._fit(
+            X, xp, is_array_api_compliant, y, sample_weight, queue=queue
+        )
         self._save_attributes()
 
     def _onedal_supported(self, method_name, *data):
@@ -178,9 +182,11 @@ def fit(self, X, y=None, sample_weight=None):
             if self.eps <= 0.0:
                 raise ValueError(f"eps == {self.eps}, must be > 0.0.")
 
+        # TODO:
+        # should be checked for Array API inputs.
         if sample_weight is not None:
             sample_weight = _check_sample_weight(sample_weight, X)
-        dispatch(
+        dispatch_with_array_api(
             self,
             "fit",
             {
@@ -194,4 +200,7 @@ def fit(self, X, y=None, sample_weight=None):
 
         return self
 
+    def _more_tags(self):
+        return {"array_api_support_sklearnex": True}
+
     fit.__doc__ = sklearn_DBSCAN.fit.__doc__

sklearnex/dispatcher.py

@@ -128,6 +128,11 @@ def get_patch_map_core(preview=False):
         from ._config import get_config as get_config_sklearnex
         from ._config import set_config as set_config_sklearnex
 
+        # TODO:
+        # check the version of skl.
+        if sklearn_check_version("1.4"):
+            import sklearn.utils._array_api as _array_api_module
+
         if sklearn_check_version("1.2.1"):
             from .utils.parallel import _FuncWrapper as _FuncWrapper_sklearnex
         else:
@@ -165,6 +170,12 @@ def get_patch_map_core(preview=False):
         from .svm import NuSVC as NuSVC_sklearnex
         from .svm import NuSVR as NuSVR_sklearnex
 
+        # TODO:
+        # check the version of skl.
+        if sklearn_check_version("1.4"):
+            from .utils._array_api import _convert_to_numpy as _convert_to_numpy_sklearnex
+            from .utils._array_api import get_namespace as get_namespace_sklearnex
+
         # DBSCAN
         mapping.pop("dbscan")
         mapping["dbscan"] = [[(cluster_module, "DBSCAN", DBSCAN_sklearnex), None]]
@@ -440,6 +451,26 @@ def get_patch_map_core(preview=False):
         mapping["_funcwrapper"] = [
             [(parallel_module, "_FuncWrapper", _FuncWrapper_sklearnex), None]
         ]
+        # TODO:
+        # check the version of skl.
+        if sklearn_check_version("1.4"):
+            # Necessary for array_api support
+            mapping["get_namespace"] = [
+                [
+                    (
+                        _array_api_module,
+                        "get_namespace",
+                        get_namespace_sklearnex,
+                    ),
+                    None,
+                ]
+            ]
+            mapping["_convert_to_numpy"] = [
+                [
+                    (_array_api_module, "_convert_to_numpy", _convert_to_numpy_sklearnex),
+                    None,
+                ]
+            ]
     return mapping