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

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -68,6 +68,7 @@
     concat,
     expand_dims,
     flip,
+    permute_dims,
     repeat,
     roll,
 )
@@ -101,6 +102,7 @@
     "isdtype",
     "linspace",
     "meshgrid",
+    "permute_dims",
     "nonzero",
     "ones",
     "ones_like",

dpctl_ext/tensor/_copy_utils.py

@@ -695,7 +695,7 @@ def _make_empty_like_orderK(x, dt, usm_type, dev):
             for i in range(x.ndim)
         )
         R = R[sl]
-    return dpt.permute_dims(R, inv_perm)
+    return dpt_ext.permute_dims(R, inv_perm)
 
 
 def _empty_like_orderK(x, dt, usm_type=None, dev=None):
@@ -800,7 +800,7 @@ def _empty_like_pair_orderK(X1, X2, dt, res_shape, usm_type, dev):
             for i in range(nd1)
         )
         R = R[sl]
-    return dpt.permute_dims(R, inv_perm)
+    return dpt_ext.permute_dims(R, inv_perm)
 
 
 def _empty_like_triple_orderK(X1, X2, X3, dt, res_shape, usm_type, dev):
@@ -876,7 +876,7 @@ def _empty_like_triple_orderK(X1, X2, X3, dt, res_shape, usm_type, dev):
             for i in range(nd1)
         )
         R = R[sl]
-    return dpt.permute_dims(R, inv_perm)
+    return dpt_ext.permute_dims(R, inv_perm)
 
 
 def copy(usm_ary, /, *, order="K"):

dpctl_ext/tensor/_manipulation_functions.py

@@ -437,6 +437,43 @@ def flip(X, /, *, axis=None):
     return X[indexer]
 
 
+def permute_dims(X, /, axes):
+    """permute_dims(x, axes)
+
+    Permute the axes (dimensions) of an array; returns the permuted
+    array as a view.
+
+    Args:
+        x (usm_ndarray): input array.
+        axes (Tuple[int, ...]): tuple containing permutation of
+           `(0,1,...,N-1)` where `N` is the number of axes (dimensions)
+           of `x`.
+    Returns:
+        usm_ndarray:
+            An array with permuted axes.
+            The returned array must has the same data type as `x`,
+            is created on the same device as `x` and has the same USM allocation
+            type as `x`.
+    """
+    if not isinstance(X, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray type, got {type(X)}.")
+    axes = normalize_axis_tuple(axes, X.ndim, "axes")
+    if not X.ndim == len(axes):
+        raise ValueError(
+            "The length of the passed axes does not match "
+            "to the number of usm_ndarray dimensions."
+        )
+    newstrides = tuple(X.strides[i] for i in axes)
+    newshape = tuple(X.shape[i] for i in axes)
+    return dpt.usm_ndarray(
+        shape=newshape,
+        dtype=X.dtype,
+        buffer=X,
+        strides=newstrides,
+        offset=X._element_offset,
+    )
+
+
 def repeat(x, repeats, /, *, axis=None):
     """repeat(x, repeats, axis=None)
 

dpctl_ext/tensor/_reshape.py

@@ -37,6 +37,10 @@
     _unravel_index,
 )
 
+# TODO: revert to `import dpctl.tensor...`
+# when dpnp fully migrates dpctl/tensor
+import dpctl_ext.tensor as dpt_ext
+
 __doc__ = "Implementation module for :func:`dpctl.tensor.reshape`."
 
 
@@ -184,7 +188,7 @@ def reshape(X, /, shape, *, order="C", copy=None):
                 src=X, dst=flat_res, sycl_queue=copy_q, depends=dep_evs
             )
         else:
-            X_t = dpt.permute_dims(X, range(X.ndim - 1, -1, -1))
+            X_t = dpt_ext.permute_dims(X, range(X.ndim - 1, -1, -1))
             hev, r_e = _copy_usm_ndarray_for_reshape(
                 src=X_t, dst=flat_res, sycl_queue=copy_q, depends=dep_evs
             )

dpnp/dpnp_array.py

@@ -2283,7 +2283,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.permute_dims(self._array_obj, axes)
+            usm_res = dpt_ext.permute_dims(self._array_obj, axes)
         return dpnp_array._create_from_usm_ndarray(usm_res)
 
     def var(