ENH: Array API dispatching

onedal/utils/_array_api.py

@@ -73,19 +73,19 @@ def _get_sycl_namespace(*arrays):
     """Get namespace of sycl arrays."""
 
     # sycl support designed to work regardless of array_api_dispatch sklearn global value
-    sycl_type = {type(x): x for x in arrays if hasattr(x, "__sycl_usm_array_interface__")}
+    sua_iface = {type(x): x for x in arrays if hasattr(x, "__sycl_usm_array_interface__")}
 
-    if len(sycl_type) > 1:
-        raise ValueError(f"Multiple SYCL types for array inputs: {sycl_type}")
+    if len(sua_iface) > 1:
+        raise ValueError(f"Multiple SYCL types for array inputs: {sua_iface}")
 
-    if sycl_type:
-        (X,) = sycl_type.values()
+    if sua_iface:
+        (X,) = sua_iface.values()
 
         if hasattr(X, "__array_namespace__"):
-            return sycl_type, X.__array_namespace__(), True
+            return sua_iface, X.__array_namespace__(), True
         elif dpnp_available and isinstance(X, dpnp.ndarray):
-            return sycl_type, dpnp, False
+            return sua_iface, dpnp, False
         else:
-            raise ValueError(f"SYCL type not recognized: {sycl_type}")
+            raise ValueError(f"SYCL type not recognized: {sua_iface}")
 
-    return sycl_type, None, False
+    return sua_iface, None, False

sklearnex/_device_offload.py

@@ -16,13 +16,14 @@
 
 from functools import wraps
 
+from daal4py.sklearn._utils import sklearn_check_version
 from onedal._device_offload import (
     _copy_to_usm,
     _get_global_queue,
     _transfer_to_host,
     dpnp_available,
 )
-from onedal.utils._array_api import _asarray, _is_numpy_namespace
+from onedal.utils._array_api import _asarray
 
 if dpnp_available:
     import dpnp
@@ -74,7 +75,7 @@ def dispatch(obj, method_name, branches, *args, **kwargs):
         return branches[backend](obj, *hostargs, **hostkwargs, queue=q)
     if backend == "sklearn":
         if (
-            "array_api_dispatch" in get_config()
+            sklearn_check_version("1.2")
             and get_config()["array_api_dispatch"]
             and "array_api_support" in obj._get_tags()
             and obj._get_tags()["array_api_support"]

sklearnex/dispatcher.py

@@ -128,6 +128,9 @@ 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
 
+        if sklearn_check_version("1.2"):
+            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 +168,10 @@ def get_patch_map_core(preview=False):
         from .svm import NuSVC as NuSVC_sklearnex
         from .svm import NuSVR as NuSVR_sklearnex
 
+        if sklearn_check_version("1.2"):
+            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 +447,24 @@ def get_patch_map_core(preview=False):
         mapping["_funcwrapper"] = [
             [(parallel_module, "_FuncWrapper", _FuncWrapper_sklearnex), None]
         ]
+        if sklearn_check_version("1.2"):
+            # 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
 
 

sklearnex/tests/test_memory_usage.py

@@ -45,6 +45,7 @@
 
 
 CPU_SKIP_LIST = (
+    "_convert_to_numpy",  # additional memory allocation is expected proportional to the input data
     "TSNE",  # too slow for using in testing on common data size
     "config_context",  # does not malloc
     "get_config",  # does not malloc
@@ -59,6 +60,7 @@
 )
 
 GPU_SKIP_LIST = (
+    "_convert_to_numpy",  # additional memory allocation is expected proportional to the input data
     "TSNE",  # too slow for using in testing on common data size
     "RandomForestRegressor",  # too slow for using in testing on common data size
     "KMeans",  # does not support GPU offloading

sklearnex/utils/_array_api.py

@@ -19,64 +19,101 @@
 import numpy as np
 
 from daal4py.sklearn._utils import sklearn_check_version
-from onedal.utils._array_api import _get_sycl_namespace
+from onedal.utils._array_api import _asarray, _get_sycl_namespace
 
+# TODO:
+# check the version of skl.
 if sklearn_check_version("1.2"):
     from sklearn.utils._array_api import get_namespace as sklearn_get_namespace
+    from sklearn.utils._array_api import _convert_to_numpy as _sklearn_convert_to_numpy
 
+from onedal._device_offload import dpctl_available, dpnp_available
 
-def get_namespace(*arrays):
-    """Get namespace of arrays.
+if dpctl_available:
+    import dpctl.tensor as dpt
 
-    Introspect `arrays` arguments and return their common Array API
-    compatible namespace object, if any. NumPy 1.22 and later can
-    construct such containers using the `numpy.array_api` namespace
-    for instance.
+if dpnp_available:
+    import dpnp
 
-    This function will return the namespace of SYCL-related arrays
-    which define the __sycl_usm_array_interface__ attribute
-    regardless of array_api support, the configuration of
-    array_api_dispatch, or scikit-learn version.
 
-    See: https://numpy.org/neps/nep-0047-array-api-standard.html
+def _convert_to_numpy(array, xp):
+    """Convert X into a NumPy ndarray on the CPU."""
+    xp_name = xp.__name__
 
-    If `arrays` are regular numpy arrays, an instance of the
-    `_NumPyApiWrapper` compatibility wrapper is returned instead.
+    if dpctl_available and xp_name in {
+        "dpctl.tensor",
+    }:
+        return dpt.to_numpy(array)
+    elif dpnp_available and isinstance(array, dpnp.ndarray):
+        return dpnp.asnumpy(array)
+    elif sklearn_check_version("1.2"):
+        return _sklearn_convert_to_numpy(array, xp)
+    else:
+        return _asarray(array, xp)
 
-    Namespace support is not enabled by default. To enabled it
-    call:
 
-      sklearn.set_config(array_api_dispatch=True)
+if sklearn_check_version("1.5"):
 
-    or:
+    def get_namespace(*arrays, remove_none=True, remove_types=(str,), xp=None):
+        """Get namespace of arrays.
 
-      with sklearn.config_context(array_api_dispatch=True):
-          # your code here
+        TBD
 
-    Otherwise an instance of the `_NumPyApiWrapper`
-    compatibility wrapper is always returned irrespective of
-    the fact that arrays implement the `__array_namespace__`
-    protocol or not.
+        Parameters
+        ----------
+        *arrays : array objects
+            Array objects.
 
-    Parameters
-    ----------
-    *arrays : array objects
-        Array objects.
+        Returns
+        -------
+        namespace : module
+            Namespace shared by 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.
+        """
 
-    is_array_api : bool
-        True of the arrays are containers that implement the Array API spec.
-    """
+        usm_iface, xp_sycl_namespace, is_array_api_compliant = _get_sycl_namespace(
+            *arrays
+        )
 
-    sycl_type, xp, is_array_api_compliant = _get_sycl_namespace(*arrays)
+        if usm_iface:
+            return xp_sycl_namespace, is_array_api_compliant
+        elif sklearn_check_version("1.2"):
+            return sklearn_get_namespace(
+                *arrays, remove_none=remove_none, remove_types=remove_types, xp=xp
+            )
+        else:
+            return np, False
 
-    if sycl_type:
-        return xp, is_array_api_compliant
-    elif sklearn_check_version("1.2"):
-        return sklearn_get_namespace(*arrays)
-    else:
-        return np, False
+else:
+
+    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.
+        """
+
+        usm_iface, xp_sycl_namespace, is_array_api_compliant = _get_sycl_namespace(
+            *arrays
+        )
+
+        if usm_iface:
+            return xp_sycl_namespace, is_array_api_compliant
+        elif sklearn_check_version("1.2"):
+            return sklearn_get_namespace(*arrays)
+        else:
+            return np, False