Switch fully to dpctl_ext.tensor in dpctl_ext.tensor

Author: vlad-perevezentsev

dpctl_ext/tensor/_accumulation.py

@@ -27,12 +27,11 @@
 # *****************************************************************************
 
 import dpctl
-import dpctl.tensor as dpt
 from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
 
 # 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_accumulation_impl as tai
 import dpctl_ext.tensor._tensor_impl as ti
 
@@ -82,7 +81,7 @@ def _accumulate_common(
         perm = [i for i in range(nd) if i != axis] + [
             axis,
         ]
-        arr = dpt_ext.permute_dims(x, perm)
+        arr = dpt.permute_dims(x, perm)
     q = x.sycl_queue
     inp_dt = x.dtype
     res_usm_type = x.usm_type
@@ -130,16 +129,16 @@ def _accumulate_common(
             )
         # permute out array dims if necessary
         if a1 != nd:
-            out = dpt_ext.permute_dims(out, perm)
+            out = dpt.permute_dims(out, perm)
             orig_out = out
         if ti._array_overlap(x, out) and implemented_types:
-            out = dpt_ext.empty_like(out)
+            out = dpt.empty_like(out)
     else:
-        out = dpt_ext.empty(
+        out = dpt.empty(
             res_sh, dtype=res_dt, usm_type=res_usm_type, sycl_queue=q
         )
         if a1 != nd:
-            out = dpt_ext.permute_dims(out, perm)
+            out = dpt.permute_dims(out, perm)
 
     _manager = SequentialOrderManager[q]
     depends = _manager.submitted_events
@@ -166,7 +165,7 @@ def _accumulate_common(
             out = orig_out
     else:
         if _dtype_supported(res_dt, res_dt):
-            tmp = dpt_ext.empty(
+            tmp = dpt.empty(
                 arr.shape, dtype=res_dt, usm_type=res_usm_type, sycl_queue=q
             )
             ht_e_cpy, cpy_e = ti._copy_usm_ndarray_into_usm_ndarray(
@@ -191,18 +190,18 @@ def _accumulate_common(
             _manager.add_event_pair(ht_e, acc_ev)
         else:
             buf_dt = _default_accumulation_type_fn(inp_dt, q)
-            tmp = dpt_ext.empty(
+            tmp = dpt.empty(
                 arr.shape, dtype=buf_dt, usm_type=res_usm_type, sycl_queue=q
             )
             ht_e_cpy, cpy_e = ti._copy_usm_ndarray_into_usm_ndarray(
                 src=arr, dst=tmp, sycl_queue=q, depends=depends
             )
             _manager.add_event_pair(ht_e_cpy, cpy_e)
-            tmp_res = dpt_ext.empty(
+            tmp_res = dpt.empty(
                 res_sh, dtype=buf_dt, usm_type=res_usm_type, sycl_queue=q
             )
             if a1 != nd:
-                tmp_res = dpt_ext.permute_dims(tmp_res, perm)
+                tmp_res = dpt.permute_dims(tmp_res, perm)
             if not include_initial:
                 ht_e, acc_ev = _accumulate_fn(
                     src=tmp,
@@ -225,10 +224,10 @@ def _accumulate_common(
             _manager.add_event_pair(ht_e_cpy2, cpy_e2)
 
     if appended_axis:
-        out = dpt_ext.squeeze(out)
+        out = dpt.squeeze(out)
     if a1 != nd:
         inv_perm = sorted(range(nd), key=lambda d: perm[d])
-        out = dpt_ext.permute_dims(out, inv_perm)
+        out = dpt.permute_dims(out, inv_perm)
 
     return out
 

dpctl_ext/tensor/_clip.py

@@ -27,12 +27,11 @@
 # *****************************************************************************
 
 import dpctl
-import dpctl.tensor as dpt
 from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
 
 # 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_elementwise_impl as tei
 import dpctl_ext.tensor._tensor_impl as ti
 
@@ -163,20 +162,20 @@ def _clip_none(x, val, out, order, _binary_fn):
 
         if ti._array_overlap(x, out):
             if not ti._same_logical_tensors(x, out):
-                out = dpt_ext.empty_like(out)
+                out = dpt.empty_like(out)
 
         if isinstance(val, dpt.usm_ndarray):
             if (
                 ti._array_overlap(val, out)
                 and not ti._same_logical_tensors(val, out)
                 and val_dtype == res_dt
             ):
-                out = dpt_ext.empty_like(out)
+                out = dpt.empty_like(out)
 
     if isinstance(val, dpt.usm_ndarray):
         val_ary = val
     else:
-        val_ary = dpt_ext.asarray(val, dtype=val_dtype, sycl_queue=exec_q)
+        val_ary = dpt.asarray(val, dtype=val_dtype, sycl_queue=exec_q)
 
     if order == "A":
         order = (
@@ -197,17 +196,17 @@ def _clip_none(x, val, out, order, _binary_fn):
                     x, val_ary, res_dt, res_shape, res_usm_type, exec_q
                 )
             else:
-                out = dpt_ext.empty(
+                out = dpt.empty(
                     res_shape,
                     dtype=res_dt,
                     usm_type=res_usm_type,
                     sycl_queue=exec_q,
                     order=order,
                 )
         if x_shape != res_shape:
-            x = dpt_ext.broadcast_to(x, res_shape)
+            x = dpt.broadcast_to(x, res_shape)
         if val_ary.shape != res_shape:
-            val_ary = dpt_ext.broadcast_to(val_ary, res_shape)
+            val_ary = dpt.broadcast_to(val_ary, res_shape)
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_binary_ev, binary_ev = _binary_fn(
@@ -229,7 +228,7 @@ def _clip_none(x, val, out, order, _binary_fn):
         if order == "K":
             buf = _empty_like_orderK(val_ary, res_dt)
         else:
-            buf = dpt_ext.empty_like(val_ary, dtype=res_dt, order=order)
+            buf = dpt.empty_like(val_ary, dtype=res_dt, order=order)
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
@@ -242,7 +241,7 @@ def _clip_none(x, val, out, order, _binary_fn):
                     x, buf, res_dt, res_shape, res_usm_type, exec_q
                 )
             else:
-                out = dpt_ext.empty(
+                out = dpt.empty(
                     res_shape,
                     dtype=res_dt,
                     usm_type=res_usm_type,
@@ -251,8 +250,8 @@ def _clip_none(x, val, out, order, _binary_fn):
                 )
 
         if x_shape != res_shape:
-            x = dpt_ext.broadcast_to(x, res_shape)
-        buf = dpt_ext.broadcast_to(buf, res_shape)
+            x = dpt.broadcast_to(x, res_shape)
+        buf = dpt.broadcast_to(buf, res_shape)
         ht_binary_ev, binary_ev = _binary_fn(
             src1=x,
             src2=buf,
@@ -313,9 +312,9 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
     if order not in ["K", "C", "F", "A"]:
         order = "K"
     if x.dtype.kind in "iu":
-        if isinstance(min, int) and min <= dpt_ext.iinfo(x.dtype).min:
+        if isinstance(min, int) and min <= dpt.iinfo(x.dtype).min:
             min = None
-        if isinstance(max, int) and max >= dpt_ext.iinfo(x.dtype).max:
+        if isinstance(max, int) and max >= dpt.iinfo(x.dtype).max:
             max = None
     if min is None and max is None:
         exec_q = x.sycl_queue
@@ -353,14 +352,14 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
 
             if ti._array_overlap(x, out):
                 if not ti._same_logical_tensors(x, out):
-                    out = dpt_ext.empty_like(out)
+                    out = dpt.empty_like(out)
                 else:
                     return out
         else:
             if order == "K":
                 out = _empty_like_orderK(x, x.dtype)
             else:
-                out = dpt_ext.empty_like(x, order=order)
+                out = dpt.empty_like(x, order=order)
 
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
@@ -519,32 +518,32 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
 
             if ti._array_overlap(x, out):
                 if not ti._same_logical_tensors(x, out):
-                    out = dpt_ext.empty_like(out)
+                    out = dpt.empty_like(out)
 
             if isinstance(min, dpt.usm_ndarray):
                 if (
                     ti._array_overlap(min, out)
                     and not ti._same_logical_tensors(min, out)
                     and buf1_dt is None
                 ):
-                    out = dpt_ext.empty_like(out)
+                    out = dpt.empty_like(out)
 
             if isinstance(max, dpt.usm_ndarray):
                 if (
                     ti._array_overlap(max, out)
                     and not ti._same_logical_tensors(max, out)
                     and buf2_dt is None
                 ):
-                    out = dpt_ext.empty_like(out)
+                    out = dpt.empty_like(out)
 
         if isinstance(min, dpt.usm_ndarray):
             a_min = min
         else:
-            a_min = dpt_ext.asarray(min, dtype=min_dtype, sycl_queue=exec_q)
+            a_min = dpt.asarray(min, dtype=min_dtype, sycl_queue=exec_q)
         if isinstance(max, dpt.usm_ndarray):
             a_max = max
         else:
-            a_max = dpt_ext.asarray(max, dtype=max_dtype, sycl_queue=exec_q)
+            a_max = dpt.asarray(max, dtype=max_dtype, sycl_queue=exec_q)
 
         if order == "A":
             order = (
@@ -572,19 +571,19 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         exec_q,
                     )
                 else:
-                    out = dpt_ext.empty(
+                    out = dpt.empty(
                         res_shape,
                         dtype=res_dt,
                         usm_type=res_usm_type,
                         sycl_queue=exec_q,
                         order=order,
                     )
             if x_shape != res_shape:
-                x = dpt_ext.broadcast_to(x, res_shape)
+                x = dpt.broadcast_to(x, res_shape)
             if a_min.shape != res_shape:
-                a_min = dpt_ext.broadcast_to(a_min, res_shape)
+                a_min = dpt.broadcast_to(a_min, res_shape)
             if a_max.shape != res_shape:
-                a_max = dpt_ext.broadcast_to(a_max, res_shape)
+                a_max = dpt.broadcast_to(a_max, res_shape)
             _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_binary_ev, binary_ev = ti._clip(
@@ -612,7 +611,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
             if order == "K":
                 buf2 = _empty_like_orderK(a_max, buf2_dt)
             else:
-                buf2 = dpt_ext.empty_like(a_max, dtype=buf2_dt, order=order)
+                buf2 = dpt.empty_like(a_max, dtype=buf2_dt, order=order)
             _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
@@ -631,18 +630,18 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         exec_q,
                     )
                 else:
-                    out = dpt_ext.empty(
+                    out = dpt.empty(
                         res_shape,
                         dtype=res_dt,
                         usm_type=res_usm_type,
                         sycl_queue=exec_q,
                         order=order,
                     )
 
-            x = dpt_ext.broadcast_to(x, res_shape)
+            x = dpt.broadcast_to(x, res_shape)
             if a_min.shape != res_shape:
-                a_min = dpt_ext.broadcast_to(a_min, res_shape)
-            buf2 = dpt_ext.broadcast_to(buf2, res_shape)
+                a_min = dpt.broadcast_to(a_min, res_shape)
+            buf2 = dpt.broadcast_to(buf2, res_shape)
             ht_binary_ev, binary_ev = ti._clip(
                 src=x,
                 min=a_min,
@@ -668,7 +667,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
             if order == "K":
                 buf1 = _empty_like_orderK(a_min, buf1_dt)
             else:
-                buf1 = dpt_ext.empty_like(a_min, dtype=buf1_dt, order=order)
+                buf1 = dpt.empty_like(a_min, dtype=buf1_dt, order=order)
             _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
@@ -687,18 +686,18 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         exec_q,
                     )
                 else:
-                    out = dpt_ext.empty(
+                    out = dpt.empty(
                         res_shape,
                         dtype=res_dt,
                         usm_type=res_usm_type,
                         sycl_queue=exec_q,
                         order=order,
                     )
 
-            x = dpt_ext.broadcast_to(x, res_shape)
-            buf1 = dpt_ext.broadcast_to(buf1, res_shape)
+            x = dpt.broadcast_to(x, res_shape)
+            buf1 = dpt.broadcast_to(buf1, res_shape)
             if a_max.shape != res_shape:
-                a_max = dpt_ext.broadcast_to(a_max, res_shape)
+                a_max = dpt.broadcast_to(a_max, res_shape)
             ht_binary_ev, binary_ev = ti._clip(
                 src=x,
                 min=buf1,
@@ -736,7 +735,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
         if order == "K":
             buf1 = _empty_like_orderK(a_min, buf1_dt)
         else:
-            buf1 = dpt_ext.empty_like(a_min, dtype=buf1_dt, order=order)
+            buf1 = dpt.empty_like(a_min, dtype=buf1_dt, order=order)
 
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
@@ -747,7 +746,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
         if order == "K":
             buf2 = _empty_like_orderK(a_max, buf2_dt)
         else:
-            buf2 = dpt_ext.empty_like(a_max, dtype=buf2_dt, order=order)
+            buf2 = dpt.empty_like(a_max, dtype=buf2_dt, order=order)
         ht_copy2_ev, copy2_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=a_max, dst=buf2, sycl_queue=exec_q, depends=dep_evs
         )
@@ -758,17 +757,17 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                     x, buf1, buf2, res_dt, res_shape, res_usm_type, exec_q
                 )
             else:
-                out = dpt_ext.empty(
+                out = dpt.empty(
                     res_shape,
                     dtype=res_dt,
                     usm_type=res_usm_type,
                     sycl_queue=exec_q,
                     order=order,
                 )
 
-        x = dpt_ext.broadcast_to(x, res_shape)
-        buf1 = dpt_ext.broadcast_to(buf1, res_shape)
-        buf2 = dpt_ext.broadcast_to(buf2, res_shape)
+        x = dpt.broadcast_to(x, res_shape)
+        buf1 = dpt.broadcast_to(buf1, res_shape)
+        buf2 = dpt.broadcast_to(buf2, res_shape)
         ht_, clip_ev = ti._clip(
             src=x,
             min=buf1,