Switch fully to dpctl_ext.tensor in dpnp

Author: vlad-perevezentsev

dpnp/__init__.py

@@ -64,7 +64,7 @@
 # Borrowed from DPCTL
 with warnings.catch_warnings():
     warnings.simplefilter("ignore", DeprecationWarning)
-    from dpctl.tensor import __array_api_version__, DLDeviceType
+    from dpctl_ext.tensor import __array_api_version__, DLDeviceType
 
 from .dpnp_array import dpnp_array as ndarray
 from .dpnp_array_api_info import __array_namespace_info__

dpnp/dpnp_algo/dpnp_arraycreation.py

@@ -29,13 +29,12 @@
 import math
 import operator
 
-import dpctl.tensor as dpt
 import dpctl.utils as dpu
 import numpy
 
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
-import dpctl_ext.tensor as dpt_ext
+import dpctl_ext.tensor as dpt
 import dpnp
 from dpnp.dpnp_array import dpnp_array
 from dpnp.dpnp_utils import get_usm_allocations, map_dtype_to_device
@@ -53,7 +52,7 @@ def _as_usm_ndarray(a, usm_type, sycl_queue):
 
     if isinstance(a, dpnp_array):
         a = a.get_array()
-    return dpt_ext.asarray(a, usm_type=usm_type, sycl_queue=sycl_queue)
+    return dpt.asarray(a, usm_type=usm_type, sycl_queue=sycl_queue)
 
 
 def _check_has_zero_val(a):
@@ -196,7 +195,7 @@ def dpnp_linspace(
 
     if dpnp.isscalar(start) and dpnp.isscalar(stop):
         # Call linspace() function for scalars.
-        usm_res = dpt_ext.linspace(
+        usm_res = dpt.linspace(
             start,
             stop,
             num,
@@ -213,37 +212,35 @@ def dpnp_linspace(
             else:
                 step = dpnp.nan
     else:
-        usm_start = dpt_ext.asarray(
+        usm_start = dpt.asarray(
             start,
             dtype=dt,
             usm_type=_usm_type,
             sycl_queue=sycl_queue_normalized,
         )
-        usm_stop = dpt_ext.asarray(
+        usm_stop = dpt.asarray(
             stop, dtype=dt, usm_type=_usm_type, sycl_queue=sycl_queue_normalized
         )
 
         delta = usm_stop - usm_start
 
-        usm_res = dpt_ext.arange(
+        usm_res = dpt.arange(
             0,
             stop=num,
             step=1,
             dtype=dt,
             usm_type=_usm_type,
             sycl_queue=sycl_queue_normalized,
         )
-        usm_res = dpt_ext.reshape(
-            usm_res, (-1,) + (1,) * delta.ndim, copy=False
-        )
+        usm_res = dpt.reshape(usm_res, (-1,) + (1,) * delta.ndim, copy=False)
 
         if step_num > 0:
             step = delta / step_num
 
             # Needed a special handling for denormal numbers (when step == 0),
             # see numpy#5437 for more details.
             # Note, dpt.where() is used to avoid a synchronization branch.
-            usm_res = dpt_ext.where(
+            usm_res = dpt.where(
                 step == 0, (usm_res / step_num) * delta, usm_res * step
             )
         else:
@@ -256,17 +253,17 @@ def dpnp_linspace(
             usm_res[-1, ...] = usm_stop
 
     if axis != 0:
-        usm_res = dpt_ext.moveaxis(usm_res, 0, axis)
+        usm_res = dpt.moveaxis(usm_res, 0, axis)
 
     if dpnp.issubdtype(dtype, dpnp.integer):
         dpt.floor(usm_res, out=usm_res)
 
-    res = dpt_ext.astype(usm_res, dtype, copy=False)
+    res = dpt.astype(usm_res, dtype, copy=False)
     res = dpnp_array._create_from_usm_ndarray(res)
 
     if retstep is True:
         if dpnp.isscalar(step):
-            step = dpt_ext.asarray(
+            step = dpt.asarray(
                 step, usm_type=res.usm_type, sycl_queue=res.sycl_queue
             )
         return res, dpnp_array._create_from_usm_ndarray(step)

dpnp/dpnp_algo/dpnp_elementwise_common.py

@@ -29,28 +29,27 @@
 import warnings
 from functools import wraps
 
-import dpctl.tensor as dpt
-import dpctl.tensor._type_utils as dtu
 import dpctl.utils as dpu
 import numpy
-from dpctl.tensor._elementwise_common import (
-    BinaryElementwiseFunc,
-    UnaryElementwiseFunc,
-)
-from dpctl.tensor._scalar_utils import (
-    _get_dtype,
-    _get_shape,
-    _validate_dtype,
-)
 
 # pylint: disable=no-name-in-module
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
-import dpctl_ext.tensor as dpt_ext
+import dpctl_ext.tensor as dpt
 import dpctl_ext.tensor._copy_utils as dtc
 import dpctl_ext.tensor._tensor_impl as dti
+import dpctl_ext.tensor._type_utils as dtu
 import dpnp
 import dpnp.backend.extensions.vm._vm_impl as vmi
+from dpctl_ext.tensor._elementwise_common import (
+    BinaryElementwiseFunc,
+    UnaryElementwiseFunc,
+)
+from dpctl_ext.tensor._scalar_utils import (
+    _get_dtype,
+    _get_shape,
+    _validate_dtype,
+)
 from dpnp.dpnp_array import dpnp_array
 from dpnp.dpnp_utils import get_usm_allocations
 from dpnp.dpnp_utils.dpnp_utils_common import (
@@ -213,7 +212,7 @@ def __call__(
 
         x_usm = dpnp.get_usm_ndarray(x)
         if dtype is not None:
-            x_usm = dpt_ext.astype(x_usm, dtype, copy=False)
+            x_usm = dpt.astype(x_usm, dtype, copy=False)
 
         out = self._unpack_out_kw(out)
         out_usm = None if out is None else dpnp.get_usm_ndarray(out)
@@ -467,7 +466,7 @@ def __call__(
                     )
 
                 # Allocate a temporary buffer with the required dtype
-                out[i] = dpt_ext.empty_like(res, dtype=res_dt)
+                out[i] = dpt.empty_like(res, dtype=res_dt)
             elif (
                 buf_dt is None
                 and dti._array_overlap(x, res)
@@ -476,7 +475,7 @@ def __call__(
                 # Allocate a temporary buffer to avoid memory overlapping.
                 # Note if `buf_dt` is not None, a temporary copy of `x` will be
                 # created, so the array overlap check isn't needed.
-                out[i] = dpt_ext.empty_like(res)
+                out[i] = dpt.empty_like(res)
 
         _manager = dpu.SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
@@ -486,7 +485,7 @@ def __call__(
             if order == "K":
                 buf = dtc._empty_like_orderK(x, buf_dt)
             else:
-                buf = dpt_ext.empty_like(x, dtype=buf_dt, order=order)
+                buf = dpt.empty_like(x, dtype=buf_dt, order=order)
 
             ht_copy_ev, copy_ev = dti._copy_usm_ndarray_into_usm_ndarray(
                 src=x, dst=buf, sycl_queue=exec_q, depends=dep_evs
@@ -503,7 +502,7 @@ def __call__(
                 if order == "K":
                     out[i] = dtc._empty_like_orderK(x, res_dt)
                 else:
-                    out[i] = dpt_ext.empty_like(x, dtype=res_dt, order=order)
+                    out[i] = dpt.empty_like(x, dtype=res_dt, order=order)
 
         # Call the unary function with input and output arrays
         ht_unary_ev, unary_ev = self.get_implementation_function()(
@@ -713,24 +712,24 @@ def __call__(
 
         if dtype is not None:
             if dpnp.isscalar(x1):
-                x1_usm = dpt_ext.asarray(
+                x1_usm = dpt.asarray(
                     x1,
                     dtype=dtype,
                     sycl_queue=x2.sycl_queue,
                     usm_type=x2.usm_type,
                 )
-                x2_usm = dpt_ext.astype(x2_usm, dtype, copy=False)
+                x2_usm = dpt.astype(x2_usm, dtype, copy=False)
             elif dpnp.isscalar(x2):
-                x1_usm = dpt_ext.astype(x1_usm, dtype, copy=False)
-                x2_usm = dpt_ext.asarray(
+                x1_usm = dpt.astype(x1_usm, dtype, copy=False)
+                x2_usm = dpt.asarray(
                     x2,
                     dtype=dtype,
                     sycl_queue=x1.sycl_queue,
                     usm_type=x1.usm_type,
                 )
             else:
-                x1_usm = dpt_ext.astype(x1_usm, dtype, copy=False)
-                x2_usm = dpt_ext.astype(x2_usm, dtype, copy=False)
+                x1_usm = dpt.astype(x1_usm, dtype, copy=False)
+                x2_usm = dpt.astype(x2_usm, dtype, copy=False)
 
         res_usm = super().__call__(x1_usm, x2_usm, out=out_usm, order=order)
 
@@ -1078,7 +1077,7 @@ def __call__(
                     )
 
                 # Allocate a temporary buffer with the required dtype
-                out[i] = dpt_ext.empty_like(res, dtype=res_dt)
+                out[i] = dpt.empty_like(res, dtype=res_dt)
             else:
                 # If `dt` is not None, a temporary copy of `x` will be created,
                 # so the array overlap check isn't needed.
@@ -1094,7 +1093,7 @@ def __call__(
                     for x in x_to_check
                 ):
                     # allocate a temporary buffer to avoid memory overlapping
-                    out[i] = dpt_ext.empty_like(res)
+                    out[i] = dpt.empty_like(res)
 
         x1 = dpnp.as_usm_ndarray(x1, dtype=x1_dt, sycl_queue=exec_q)
         x2 = dpnp.as_usm_ndarray(x2, dtype=x2_dt, sycl_queue=exec_q)
@@ -1127,7 +1126,7 @@ def __call__(
                 if order == "K":
                     buf = dtc._empty_like_orderK(x, buf_dt)
                 else:
-                    buf = dpt_ext.empty_like(x, dtype=buf_dt, order=order)
+                    buf = dpt.empty_like(x, dtype=buf_dt, order=order)
 
                 ht_copy_ev, copy_ev = dti._copy_usm_ndarray_into_usm_ndarray(
                     src=x, dst=buf, sycl_queue=exec_q, depends=dep_evs
@@ -1146,7 +1145,7 @@ def __call__(
                         x1, x2, res_dt, res_shape, res_usm_type, exec_q
                     )
                 else:
-                    out[i] = dpt_ext.empty(
+                    out[i] = dpt.empty(
                         res_shape,
                         dtype=res_dt,
                         order=order,
@@ -1156,9 +1155,9 @@ def __call__(
 
         # Broadcast shapes of input arrays
         if x1.shape != res_shape:
-            x1 = dpt_ext.broadcast_to(x1, res_shape)
+            x1 = dpt.broadcast_to(x1, res_shape)
         if x2.shape != res_shape:
-            x2 = dpt_ext.broadcast_to(x2, res_shape)
+            x2 = dpt.broadcast_to(x2, res_shape)
 
         # Call the binary function with input and output arrays
         ht_binary_ev, binary_ev = self.get_implementation_function()(
@@ -1326,7 +1325,7 @@ def __call__(self, x, /, decimals=0, out=None, *, dtype=None):
                 res_usm = dpt.divide(x_usm, 10**decimals, out=out_usm)
 
             if dtype is not None:
-                res_usm = dpt_ext.astype(res_usm, dtype, copy=False)
+                res_usm = dpt.astype(res_usm, dtype, copy=False)
 
             if out is not None and isinstance(out, dpnp_array):
                 return out

dpnp/dpnp_array.py

@@ -37,11 +37,9 @@
 
 import warnings
 
-import dpctl.tensor as dpt
-
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
-import dpctl_ext.tensor as dpt_ext
+import dpctl_ext.tensor as dpt
 import dpctl_ext.tensor._type_utils as dtu
 import dpnp
 from dpctl_ext.tensor._numpy_helper import AxisError
@@ -777,7 +775,7 @@ def asnumpy(self):
 
         """
 
-        return dpt_ext.asnumpy(self._array_obj)
+        return dpt.asnumpy(self._array_obj)
 
     def astype(
         self,
@@ -2283,7 +2281,7 @@ def transpose(self, *axes):
                 # self.transpose(None).shape == self.shape[::-1]
                 axes = tuple((ndim - x - 1) for x in range(ndim))
 
-            usm_res = dpt_ext.permute_dims(self._array_obj, axes)
+            usm_res = dpt.permute_dims(self._array_obj, axes)
         return dpnp_array._create_from_usm_ndarray(usm_res)
 
     def var(

dpnp/dpnp_array_api_info.py

@@ -36,7 +36,9 @@
 
 """
 
-import dpctl.tensor as dpt
+# TODO: revert to `import dpctl.tensor...`
+# when dpnp fully migrates dpctl/tensor
+import dpctl_ext.tensor as dpt
 
 
 def __array_namespace_info__():

dpnp/dpnp_iface.py

@@ -45,17 +45,16 @@
 import os
 
 import dpctl
-import dpctl.tensor as dpt
 import dpctl.utils as dpu
 import numpy
-from dpctl.tensor._device import normalize_queue_device
 
 # pylint: disable=no-name-in-module
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
-import dpctl_ext.tensor as dpt_ext
+import dpctl_ext.tensor as dpt
 import dpctl_ext.tensor._tensor_impl as ti
 import dpnp
+from dpctl_ext.tensor._device import normalize_queue_device
 
 from .dpnp_array import dpnp_array
 from .dpnp_utils import (
@@ -137,7 +136,7 @@ def asnumpy(a, order="C"):
         return a.asnumpy()
 
     if isinstance(a, dpt.usm_ndarray):
-        return dpt_ext.asnumpy(a)
+        return dpt.asnumpy(a)
 
     return numpy.asarray(a, order=order)
 
@@ -191,7 +190,7 @@ def as_usm_ndarray(a, dtype=None, device=None, usm_type=None, sycl_queue=None):
     if is_supported_array_type(a):
         return get_usm_ndarray(a)
 
-    return dpt_ext.asarray(
+    return dpt.asarray(
         a, dtype=dtype, device=device, usm_type=usm_type, sycl_queue=sycl_queue
     )