Move ti.unstack() to dpctl_ext.tensor and reuse it

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -76,6 +76,7 @@
     stack,
     swapaxes,
     tile,
+    unstack,
 )
 from dpctl_ext.tensor._reshape import reshape
 
@@ -128,6 +129,7 @@
     "to_numpy",
     "tril",
     "triu",
+    "unstack",
     "where",
     "zeros",
     "zeros_like",

dpctl_ext/tensor/_manipulation_functions.py

@@ -974,6 +974,35 @@ def swapaxes(X, axis1, axis2):
     return dpt_ext.permute_dims(X, tuple(ind))
 
 
+def unstack(X, /, *, axis=0):
+    """unstack(x, axis=0)
+
+    Splits an array in a sequence of arrays along the given axis.
+
+    Args:
+        x (usm_ndarray): input array
+
+        axis (int, optional): axis along which `x` is unstacked.
+            If `x` has rank (i.e, number of dimensions) `N`,
+            a valid `axis` must reside in the half-open interval `[-N, N)`.
+            Default: `0`.
+
+    Returns:
+        Tuple[usm_ndarray,...]:
+            Output sequence of arrays which are views into the input array.
+
+    Raises:
+        AxisError: if the `axis` value is invalid.
+    """
+    if not isinstance(X, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray type, got {type(X)}.")
+
+    axis = normalize_axis_index(axis, X.ndim)
+    Y = dpt_ext.moveaxis(X, axis, 0)
+
+    return tuple(Y[i] for i in range(Y.shape[0]))
+
+
 def tile(x, repetitions, /):
     """tile(x, repetitions)
 

dpnp/dpnp_iface_manipulation.py

@@ -1093,7 +1093,9 @@ def broadcast_arrays(*args, subok=False):
     if len(args) == 0:
         return []
 
-    usm_arrays = dpt.broadcast_arrays(*[dpnp.get_usm_ndarray(a) for a in args])
+    usm_arrays = dpt_ext.broadcast_arrays(
+        *[dpnp.get_usm_ndarray(a) for a in args]
+    )
     return [dpnp_array._create_from_usm_ndarray(a) for a in usm_arrays]
 
 
@@ -1521,7 +1523,7 @@ def copyto(dst, src, casting="same_kind", where=True):
                 f"but got {where.dtype}"
             )
 
-        dst_usm, src_usm, mask_usm = dpt.broadcast_arrays(
+        dst_usm, src_usm, mask_usm = dpt_ext.broadcast_arrays(
             dpnp.get_usm_ndarray(dst),
             dpnp.get_usm_ndarray(src),
             dpnp.get_usm_ndarray(where),
@@ -4522,7 +4524,7 @@ def unstack(x, /, *, axis=0):
     if usm_x.ndim == 0:
         raise ValueError("Input array must be at least 1-d.")
 
-    res = dpt.unstack(usm_x, axis=axis)
+    res = dpt_ext.unstack(usm_x, axis=axis)
     return tuple(dpnp_array._create_from_usm_ndarray(a) for a in res)