Extend ._tensor_impl with linear sequences functions

dpctl_ext/tensor/CMakeLists.txt

@@ -51,7 +51,7 @@ set(_tensor_impl_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/copy_numpy_ndarray_into_usm_ndarray.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/copy_for_reshape.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/copy_for_roll.cpp
-    # ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/linear_sequences.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/linear_sequences.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/integer_advanced_indexing.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/boolean_advanced_indexing.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/eye_ctor.cpp
@@ -93,7 +93,7 @@ endif()
 set(_no_fast_math_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/copy_and_cast_usm_to_usm.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/full_ctor.cpp
-    # ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/linear_sequences.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/linear_sequences.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/clip.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/where.cpp
 )

dpctl_ext/tensor/__init__.py

@@ -36,10 +36,21 @@
     to_numpy,
 )
 from ._ctors import (
+    arange,
+    asarray,
+    empty,
+    empty_like,
     eye,
     full,
+    full_like,
+    linspace,
+    meshgrid,
+    ones,
+    ones_like,
     tril,
     triu,
+    zeros,
+    zeros_like,
 )
 from ._indexing_functions import (
     extract,
@@ -51,38 +62,73 @@
     take_along_axis,
 )
 from ._manipulation_functions import (
+    broadcast_arrays,
+    broadcast_to,
+    concat,
+    expand_dims,
+    flip,
+    moveaxis,
+    permute_dims,
     repeat,
     roll,
+    squeeze,
+    stack,
+    swapaxes,
+    tile,
+    unstack,
 )
 from ._reshape import reshape
 from ._search_functions import where
 from ._type_utils import can_cast, finfo, iinfo, isdtype, result_type
 
 __all__ = [
+    "arange",
+    "asarray",
     "asnumpy",
     "astype",
+    "broadcast_arrays",
+    "broadcast_to",
     "can_cast",
+    "concat",
     "copy",
     "clip",
+    "empty",
+    "empty_like",
     "extract",
+    "expand_dims",
     "eye",
     "finfo",
+    "flip",
     "from_numpy",
     "full",
+    "full_like",
     "iinfo",
     "isdtype",
+    "linspace",
+    "meshgrid",
+    "moveaxis",
+    "permute_dims",
     "nonzero",
+    "ones",
+    "ones_like",
     "place",
     "put",
     "put_along_axis",
     "repeat",
     "reshape",
     "result_type",
     "roll",
+    "squeeze",
+    "stack",
+    "swapaxes",
     "take",
     "take_along_axis",
+    "tile",
     "to_numpy",
     "tril",
     "triu",
+    "unstack",
     "where",
+    "zeros",
+    "zeros_like",
 ]

dpctl_ext/tensor/_clip.py

@@ -163,20 +163,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.empty_like(out)
+                out = dpt_ext.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.empty_like(out)
+                out = dpt_ext.empty_like(out)
 
     if isinstance(val, dpt.usm_ndarray):
         val_ary = val
     else:
-        val_ary = dpt.asarray(val, dtype=val_dtype, sycl_queue=exec_q)
+        val_ary = dpt_ext.asarray(val, dtype=val_dtype, sycl_queue=exec_q)
 
     if order == "A":
         order = (
@@ -197,17 +197,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.empty(
+                out = dpt_ext.empty(
                     res_shape,
                     dtype=res_dt,
                     usm_type=res_usm_type,
                     sycl_queue=exec_q,
                     order=order,
                 )
         if x_shape != res_shape:
-            x = dpt.broadcast_to(x, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
         if val_ary.shape != res_shape:
-            val_ary = dpt.broadcast_to(val_ary, res_shape)
+            val_ary = dpt_ext.broadcast_to(val_ary, res_shape)
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_binary_ev, binary_ev = _binary_fn(
@@ -229,7 +229,7 @@ def _clip_none(x, val, out, order, _binary_fn):
         if order == "K":
             buf = _empty_like_orderK(val_ary, res_dt)
         else:
-            buf = dpt.empty_like(val_ary, dtype=res_dt, order=order)
+            buf = dpt_ext.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 +242,7 @@ def _clip_none(x, val, out, order, _binary_fn):
                     x, buf, res_dt, res_shape, res_usm_type, exec_q
                 )
             else:
-                out = dpt.empty(
+                out = dpt_ext.empty(
                     res_shape,
                     dtype=res_dt,
                     usm_type=res_usm_type,
@@ -251,8 +251,8 @@ def _clip_none(x, val, out, order, _binary_fn):
                 )
 
         if x_shape != res_shape:
-            x = dpt.broadcast_to(x, res_shape)
-        buf = dpt.broadcast_to(buf, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
+        buf = dpt_ext.broadcast_to(buf, res_shape)
         ht_binary_ev, binary_ev = _binary_fn(
             src1=x,
             src2=buf,
@@ -353,14 +353,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.empty_like(out)
+                    out = dpt_ext.empty_like(out)
                 else:
                     return out
         else:
             if order == "K":
                 out = _empty_like_orderK(x, x.dtype)
             else:
-                out = dpt.empty_like(x, order=order)
+                out = dpt_ext.empty_like(x, order=order)
 
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
@@ -519,32 +519,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.empty_like(out)
+                    out = dpt_ext.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.empty_like(out)
+                    out = dpt_ext.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.empty_like(out)
+                    out = dpt_ext.empty_like(out)
 
         if isinstance(min, dpt.usm_ndarray):
             a_min = min
         else:
-            a_min = dpt.asarray(min, dtype=min_dtype, sycl_queue=exec_q)
+            a_min = dpt_ext.asarray(min, dtype=min_dtype, sycl_queue=exec_q)
         if isinstance(max, dpt.usm_ndarray):
             a_max = max
         else:
-            a_max = dpt.asarray(max, dtype=max_dtype, sycl_queue=exec_q)
+            a_max = dpt_ext.asarray(max, dtype=max_dtype, sycl_queue=exec_q)
 
         if order == "A":
             order = (
@@ -572,19 +572,19 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         exec_q,
                     )
                 else:
-                    out = dpt.empty(
+                    out = dpt_ext.empty(
                         res_shape,
                         dtype=res_dt,
                         usm_type=res_usm_type,
                         sycl_queue=exec_q,
                         order=order,
                     )
             if x_shape != res_shape:
-                x = dpt.broadcast_to(x, res_shape)
+                x = dpt_ext.broadcast_to(x, res_shape)
             if a_min.shape != res_shape:
-                a_min = dpt.broadcast_to(a_min, res_shape)
+                a_min = dpt_ext.broadcast_to(a_min, res_shape)
             if a_max.shape != res_shape:
-                a_max = dpt.broadcast_to(a_max, res_shape)
+                a_max = dpt_ext.broadcast_to(a_max, res_shape)
             _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_binary_ev, binary_ev = ti._clip(
@@ -612,7 +612,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.empty_like(a_max, dtype=buf2_dt, order=order)
+                buf2 = dpt_ext.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 +631,18 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         exec_q,
                     )
                 else:
-                    out = dpt.empty(
+                    out = dpt_ext.empty(
                         res_shape,
                         dtype=res_dt,
                         usm_type=res_usm_type,
                         sycl_queue=exec_q,
                         order=order,
                     )
 
-            x = dpt.broadcast_to(x, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
             if a_min.shape != res_shape:
-                a_min = dpt.broadcast_to(a_min, res_shape)
-            buf2 = dpt.broadcast_to(buf2, res_shape)
+                a_min = dpt_ext.broadcast_to(a_min, res_shape)
+            buf2 = dpt_ext.broadcast_to(buf2, res_shape)
             ht_binary_ev, binary_ev = ti._clip(
                 src=x,
                 min=a_min,
@@ -668,7 +668,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.empty_like(a_min, dtype=buf1_dt, order=order)
+                buf1 = dpt_ext.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 +687,18 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         exec_q,
                     )
                 else:
-                    out = dpt.empty(
+                    out = dpt_ext.empty(
                         res_shape,
                         dtype=res_dt,
                         usm_type=res_usm_type,
                         sycl_queue=exec_q,
                         order=order,
                     )
 
-            x = dpt.broadcast_to(x, res_shape)
-            buf1 = dpt.broadcast_to(buf1, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
+            buf1 = dpt_ext.broadcast_to(buf1, res_shape)
             if a_max.shape != res_shape:
-                a_max = dpt.broadcast_to(a_max, res_shape)
+                a_max = dpt_ext.broadcast_to(a_max, res_shape)
             ht_binary_ev, binary_ev = ti._clip(
                 src=x,
                 min=buf1,
@@ -736,7 +736,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.empty_like(a_min, dtype=buf1_dt, order=order)
+            buf1 = dpt_ext.empty_like(a_min, dtype=buf1_dt, order=order)
 
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
@@ -747,7 +747,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.empty_like(a_max, dtype=buf2_dt, order=order)
+            buf2 = dpt_ext.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 +758,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.empty(
+                out = dpt_ext.empty(
                     res_shape,
                     dtype=res_dt,
                     usm_type=res_usm_type,
                     sycl_queue=exec_q,
                     order=order,
                 )
 
-        x = dpt.broadcast_to(x, res_shape)
-        buf1 = dpt.broadcast_to(buf1, res_shape)
-        buf2 = dpt.broadcast_to(buf2, res_shape)
+        x = dpt_ext.broadcast_to(x, res_shape)
+        buf1 = dpt_ext.broadcast_to(buf1, res_shape)
+        buf2 = dpt_ext.broadcast_to(buf2, res_shape)
         ht_, clip_ev = ti._clip(
             src=x,
             min=buf1,