[maintenance] lazy load dpnp.tensor/dpnp and prepare for array_api lazy importing

onedal/_device_offload.py

@@ -16,7 +16,6 @@
 
 import inspect
 import logging
-from collections.abc import Iterable
 from functools import wraps
 
 import numpy as np
@@ -25,14 +24,10 @@
 from onedal import _default_backend as backend
 
 from ._config import _get_config
+from .datatypes import copy_to_dpnp, copy_to_usm, usm_to_numpy
 from .utils import _sycl_queue_manager as QM
 from .utils._array_api import _asarray, _is_numpy_namespace
-from .utils._dpep_helpers import dpctl_available, dpnp_available
-
-if dpctl_available:
-    from dpctl.memory import MemoryUSMDevice, as_usm_memory
-    from dpctl.tensor import usm_ndarray
-
+from .utils._third_party import is_dpnp_ndarray
 
 logger = logging.getLogger("sklearnex")
 cpu_dlpack_device = (backend.kDLCPU, 0)
@@ -69,61 +64,13 @@
     return wrapper
 
 
-if dpnp_available:
-    import dpnp
-
-    from .utils._array_api import _convert_to_dpnp
-
-
-def _copy_to_usm(queue, array):
-    if not dpctl_available:
-        raise RuntimeError(
-            "dpctl need to be installed to work " "with __sycl_usm_array_interface__"
-        )
-
-    if hasattr(array, "__array__"):
-
-        try:
-            mem = MemoryUSMDevice(array.nbytes, queue=queue)
-            mem.copy_from_host(array.tobytes())
-            return usm_ndarray(array.shape, array.dtype, buffer=mem)
-        except ValueError as e:
-            # ValueError will raise if device does not support the dtype
-            # retry with float32 (needed for fp16 and fp64 support issues)
-            # try again as float32, if it is a float32 just raise the error.
-            if array.dtype == np.float32:
-                raise e
-            return _copy_to_usm(queue, array.astype(np.float32))
-    else:
-        if isinstance(array, Iterable):
-            array = [_copy_to_usm(queue, i) for i in array]
-        return array
-
-
 def _transfer_to_host(*data):
     has_usm_data, has_host_data = False, False
 
     host_data = []
     for item in data:
-        usm_iface = getattr(item, "__sycl_usm_array_interface__", None)
-        if usm_iface is not None:
-            if not dpctl_available:
-                raise RuntimeError(
-                    "dpctl need to be installed to work "
-                    "with __sycl_usm_array_interface__"
-                )
-
-            buffer = as_usm_memory(item).copy_to_host()
-            order = "C"
-            if usm_iface["strides"] is not None and len(usm_iface["strides"]) > 1:
-                if usm_iface["strides"][0] < usm_iface["strides"][1]:
-                    order = "F"
-            item = np.ndarray(
-                shape=usm_iface["shape"],
-                dtype=usm_iface["typestr"],
-                buffer=buffer,
-                order=order,
-            )
+        if usm_iface := getattr(item, "__sycl_usm_array_interface__", None):
+            item = usm_to_numpy(item, usm_iface)
             has_usm_data = True
         elif not isinstance(item, np.ndarray) and (
             device := getattr(item, "__dlpack_device__", None)
@@ -167,82 +114,83 @@
     _, hostvalues = _transfer_to_host(*kwargs.values())
     hostkwargs = dict(zip(kwargs.keys(), hostvalues))
     return hostargs, hostkwargs
 
 
 def support_input_format(func):
     """Transform input and output function arrays to/from host.
 
     Converts and moves the output arrays of the decorated function
     to match the input array type and device.
     Puts SYCLQueue from data to decorated function arguments.
 
     Parameters
     ----------
     func : callable
        Function or method which has array data as input.
 
     Returns
     -------
     wrapper_impl : callable
         Wrapped function or method which will return matching format.
     """
 
     def invoke_func(self_or_None, *args, **kwargs):
         if self_or_None is None:
             return func(*args, **kwargs)
         else:
             return func(self_or_None, *args, **kwargs)
 
     @wraps(func)
     def wrapper_impl(*args, **kwargs):
         # remove self from args if it is a class method
         if inspect.isfunction(func) and "." in func.__qualname__:
             self = args[0]
             args = args[1:]
         else:
             self = None
 
         # KNeighbors*.fit can not be used with raw inputs, ignore `use_raw_input=True`
         override_raw_input = (
             self
             and self.__class__.__name__ in ("KNeighborsClassifier", "KNeighborsRegressor")
             and func.__name__ == "fit"
         )
         if override_raw_input:
             pretty_name = f"{self.__class__.__name__}.{func.__name__}"
             logger.warning(
                 f"Using raw inputs is not supported for {pretty_name}. Ignoring `use_raw_input=True` setting."
             )
         if _get_config()["use_raw_input"] is True and not override_raw_input:
             if "queue" not in kwargs:
-                usm_iface = getattr(args[0], "__sycl_usm_array_interface__", None)
-                data_queue = usm_iface["syclobj"] if usm_iface is not None else None
-                kwargs["queue"] = data_queue
+                if usm_iface := getattr(args[0], "__sycl_usm_array_interface__", None):
+                    kwargs["queue"] = usm_iface["syclobj"]
+                else:
+                    kwargs["queue"] = None
             return invoke_func(self, *args, **kwargs)
         elif len(args) == 0 and len(kwargs) == 0:
             # no arguments, there's nothing we can deduce from them -> just call the function
             return invoke_func(self, *args, **kwargs)
 
-        data = (*args, *kwargs.values())
+        data = (*args, *kwargs.values())[0]
         # get and set the global queue from the kwarg or data
         with QM.manage_global_queue(kwargs.get("queue"), *args) as queue:
             hostargs, hostkwargs = _get_host_inputs(*args, **kwargs)
             if "queue" in inspect.signature(func).parameters:
                 # set the queue if it's expected by func
                 hostkwargs["queue"] = queue
             result = invoke_func(self, *hostargs, **hostkwargs)
 
-            usm_iface = getattr(data[0], "__sycl_usm_array_interface__", None)
-            if queue is not None and usm_iface is not None:
-                result = _copy_to_usm(queue, result)
-                if dpnp_available and isinstance(data[0], dpnp.ndarray):
-                    result = _convert_to_dpnp(result)
-                return result
+            if queue and hasattr(data, "__sycl_usm_array_interface__"):
+                return (
+                    copy_to_dpnp(queue, result)
+                    if is_dpnp_ndarray(data)
+                    else copy_to_usm(queue, result)
+                )
 
         if get_config().get("transform_output") in ("default", None):
-            input_array_api = getattr(data[0], "__array_namespace__", lambda: None)()
+            input_array_api = getattr(data, "__array_namespace__", lambda: None)()
             if input_array_api and not _is_numpy_namespace(input_array_api):
-                input_array_api_device = data[0].device
+                input_array_api_device = data.device
                 result = _asarray(result, input_array_api, device=input_array_api_device)
         return result
 

onedal/common/tests/test_sycl.py

@@ -18,7 +18,7 @@
 
 from onedal import _default_backend as backend
 from onedal.tests.utils._device_selection import get_queues
-from onedal.utils._dpep_helpers import dpctl_available
+from onedal.utils._third_party import dpctl_available
 
 
 @pytest.mark.skipif(

onedal/datatypes/__init__.py

@@ -15,5 +15,12 @@
 # ==============================================================================
 
 from ._data_conversion import from_table, to_table
+from ._sycl_usm import copy_to_dpnp, copy_to_usm, usm_to_numpy
 
-__all__ = ["from_table", "to_table"]
+__all__ = [
+    "copy_to_dpnp",
+    "copy_to_usm",
+    "from_table",
+    "to_table",
+    "usm_to_numpy",
+]

onedal/datatypes/_sycl_usm.py

@@ -0,0 +1,78 @@
+# ==============================================================================
+# Copyright Contributors to the oneDAL Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from collections.abc import Iterable
+
+import numpy as np
+
+from ..utils._third_party import lazy_import
+
+
+@lazy_import("dpctl.memory", "dpctl.tensor")
+def array_to_usm(memory, tensor, queue, array):
+    try:
+        mem = memory.MemoryUSMDevice(array.nbytes, queue=queue)
+        mem.copy_from_host(array.tobytes())
+        return tensor.usm_ndarray(array.shape, array.dtype, buffer=mem)
+    except ValueError as e:
+        # ValueError will raise if device does not support the dtype
+        # retry with float32 (needed for fp16 and fp64 support issues)
+        # try again as float32, if it is a float32 just raise the error.
+        if array.dtype == np.float32:
+            raise e
+        return _array_to_usm(queue, array.astype(np.float32))
+
+
+@lazy_import("dpnp", "dpctl.tensor")
+def to_dpnp(dpnp, tensor, array):
+    if isinstance(array, tensor.usm_ndarray):
+        return dpnp.array(array, copy=False)
+    else:
+        return array
+
+
+def copy_to_usm(queue, array):
+    if hasattr(array, "__array__"):
+        return array_to_usm(queue, array)
+    else:
+        if isinstance(array, Iterable):
+            array = [copy_to_usm(queue, i) for i in array]
+        return array
+
+
+def copy_to_dpnp(queue, array):
+    if hasattr(array, "__array__"):
+        return to_dpnp(array_to_usm(queue, array))
+    else:
+        if isinstance(array, Iterable):
+            array = [copy_to_dpnp(queue, i) for i in array]
+        return array
+
+
+@lazy_import("dpctl.memory")
+def usm_to_numpy(memorymod, item, usm_iface):
+    buffer = memorymod.as_usm_memory(item).copy_to_host()
+    order = "C"
+    if usm_iface["strides"] is not None and len(usm_iface["strides"]) > 1:
+        if usm_iface["strides"][0] < usm_iface["strides"][1]:
+            order = "F"
+    item = np.ndarray(
+        shape=usm_iface["shape"],
+        dtype=usm_iface["typestr"],
+        buffer=buffer,
+        order=order,
+    )
+    return item

onedal/datatypes/tests/common.py

@@ -14,11 +14,16 @@
 # limitations under the License.
 # ===============================================================================
 
-from onedal.utils._dpep_helpers import dpctl_available, dpnp_available
+from onedal.utils._third_party import dpctl_available
 
-if dpnp_available:
+try:
     import dpnp
 
+    dpnp_available = True
+except ImportError:
+    dpnp_available = False
+
+
 if dpctl_available:
     import dpctl
     from dpctl.tensor import usm_ndarray

onedal/datatypes/tests/test_data.py

@@ -24,7 +24,7 @@
 from onedal import _default_backend, _dpc_backend
 from onedal._device_offload import supports_queue
 from onedal.datatypes import from_table, to_table
-from onedal.utils._dpep_helpers import dpctl_available
+from onedal.utils._third_party import dpctl_available
 
 backend = _dpc_backend or _default_backend
 

onedal/ensemble/forest.py

@@ -34,7 +34,6 @@
 from ..common._mixin import ClassifierMixin, RegressorMixin
 from ..datatypes import from_table, to_table
 from ..utils._array_api import _get_sycl_namespace
-from ..utils._dpep_helpers import get_unique_values_with_dpep
 from ..utils.validation import (
     _check_array,
     _check_n_features,
@@ -315,7 +314,10 @@
         else:
             if sua_iface is not None:
                 queue = X.sycl_queue
-            self.classes_ = get_unique_values_with_dpep(y)
+            try:
+                self.classes_ = xp.unique(y)
+            except AttributeError:
+                self.classes_ = xp.unique_values(y)
 
         self.n_features_in_ = X.shape[1]
 
@@ -424,7 +426,7 @@
         else:
             result = self.infer(params, model, X)
 
         # TODO: fix probabilities out of [0, 1] interval on oneDAL side
         pred = from_table(result.probabilities)
         return pred.clip(0.0, 1.0)
 

onedal/linear_model/logistic_regression.py

@@ -29,7 +29,6 @@
 from ..common._mixin import ClassifierMixin
 from ..datatypes import from_table, to_table
 from ..utils._array_api import _get_sycl_namespace
-from ..utils._dpep_helpers import get_unique_values_with_dpep
 from ..utils.validation import (
     _check_array,
     _check_n_features,
@@ -96,7 +95,12 @@ def _fit(self, X, y):
             self.classes_, y = np.unique(y, return_inverse=True)
             y = y.astype(dtype=np.int32)
         else:
-            self.classes_ = get_unique_values_with_dpep(y)
+            _, xp, _ = _get_sycl_namespace(X)
+            try:
+                self.classes_ = xp.unique(y)
+            except AttributeError:
+                self.classes_ = xp.unique_values(y)
+
             n_classes = len(self.classes_)
             if n_classes != 2:
                 raise ValueError("Only binary classification is supported")

onedal/tests/utils/_dataframes_support.py

@@ -20,14 +20,18 @@
 
 from sklearnex import get_config
 
-from ...utils._dpep_helpers import dpctl_available, dpnp_available
+from ...utils._third_party import dpctl_available
 
 if dpctl_available:
     import dpctl.tensor as dpt
 
-if dpnp_available:
+try:
     import dpnp
 
+    dpnp_available = True
+except ImportError:
+    dpnp_available = False
+
 try:
     # This should be lazy imported in the
     # future along with other popular

onedal/tests/utils/_device_selection.py

@@ -18,7 +18,7 @@
 
 import pytest
 
-from ...utils._dpep_helpers import dpctl_available
+from ...utils._third_party import dpctl_available
 
 if dpctl_available:
     import dpctl