[EXPERIMENT] Enabling dpctl.tensor for scikit-learn array API conformance testing

onedal/datatypes/_data_conversion.py

@@ -32,10 +32,14 @@
     dpctl_available = False
 
 
-def _apply_and_pass(func, *args):
+def _apply_and_pass(func, *args, **kwargs):
     if len(args) == 1:
-        return func(args[0])
-    return tuple(map(func, args))
+        return func(args[0], **kwargs) if len(kwargs) > 0 else func(args[0])
+    return (
+        tuple(func(arg, **kwargs) for arg in args)
+        if len(kwargs) > 0
+        else tuple(func(arg) for arg in args)
+    )
 
 
 def from_table(*args):
@@ -59,7 +63,7 @@ def to_table(*args):
 if _is_dpc_backend:
     from ..common._policy import _HostInteropPolicy
 
-    def _convert_to_supported(policy, *data):
+    def _convert_to_supported(policy, *data, xp=np):
         def func(x):
             return x
 
@@ -71,13 +75,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 xp.astype(x, dtype=xp.float32)
             else:
                 return x
 
@@ -88,7 +92,7 @@ def convert_or_pass(x):
 
 else:
 
-    def _convert_to_supported(policy, *data):
+    def _convert_to_supported(policy, *data, xp=np):
         def func(x):
             return x
 

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.4")
             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.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 +168,14 @@ 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.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
+            from .utils._array_api import (
+                yield_namespace_device_dtype_combinations as yield_namespace_device_dtype_combinations_sklearnex,
+            )
+            from .utils._array_api import yield_namespaces as yield_namespaces_sklearnex
+
         # DBSCAN
         mapping.pop("dbscan")
         mapping["dbscan"] = [[(cluster_module, "DBSCAN", DBSCAN_sklearnex), None]]
@@ -440,6 +451,40 @@ def get_patch_map_core(preview=False):
         mapping["_funcwrapper"] = [
             [(parallel_module, "_FuncWrapper", _FuncWrapper_sklearnex), None]
         ]
+        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,
+                ]
+            ]
+            mapping["yield_namespace"] = [
+                [
+                    (_array_api_module, "yield_namespace", yield_namespaces_sklearnex),
+                    None,
+                ]
+            ]
+            mapping["yield_namespace_device_dtype_combinations"] = [
+                [
+                    (
+                        _array_api_module,
+                        "yield_namespace_device_dtype_combinations",
+                        yield_namespace_device_dtype_combinations_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/tests/test_patching.py

@@ -307,10 +307,13 @@ def list_all_attr(string):
 
     module_map = {i: i for i in sklearnex__all__.intersection(sklearn__all__)}
 
-    # _assert_all_finite patches an internal sklearn function which isn't
-    # exposed via __all__ in sklearn. It is a special case where this rule
-    # is not applied (e.g. it is grandfathered in).
+    # _assert_all_finite, _convert_to_numpy, get_namespace patch an internal
+    # sklearn functions which aren't exposed via __all__ in sklearn. It is a special
+    # case where this rule is not applied (e.g. it is grandfathered in).
     del patched["_assert_all_finite"]
+    if sklearn_check_version("1.4"):
+        del patched["_convert_to_numpy"]
+        del patched["get_namespace"]
 
     # remove all scikit-learn-intelex-only estimators
     for i in patched.copy():