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

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -66,6 +66,7 @@
     broadcast_arrays,
     broadcast_to,
     concat,
+    expand_dims,
     repeat,
     roll,
 )
@@ -88,6 +89,7 @@
     "empty",
     "empty_like",
     "extract",
+    "expand_dims",
     "eye",
     "finfo",
     "from_numpy",

dpctl_ext/tensor/_manipulation_functions.py

@@ -359,6 +359,52 @@ def concat(arrays, /, *, axis=0):
     return res
 
 
+def expand_dims(X, /, *, axis=0):
+    """expand_dims(x, axis)
+
+    Expands the shape of an array by inserting a new axis (dimension)
+    of size one at the position specified by axis.
+
+    Args:
+        x (usm_ndarray):
+            input array
+        axis (Union[int, Tuple[int]]):
+            axis position in the expanded axes (zero-based). If `x` has rank
+            (i.e, number of dimensions) `N`, a valid `axis` must reside
+            in the closed-interval `[-N-1, N]`. If provided a negative
+            `axis`, the `axis` position at which to insert a singleton
+            dimension is computed as `N + axis + 1`. Hence, if
+            provided `-1`, the resolved axis position is `N` (i.e.,
+            a singleton dimension must be appended to the input array `x`).
+            If provided `-N-1`, the resolved axis position is `0` (i.e., a
+            singleton dimension is prepended to the input array `x`).
+
+    Returns:
+        usm_ndarray:
+            Returns a view, if possible, and a copy otherwise with the number
+            of dimensions increased.
+            The expanded array has the same data type as the input array `x`.
+            The expanded array is located on the same device as the input
+            array, and has the same USM allocation type.
+
+    Raises:
+        IndexError: if `axis` value is invalid.
+    """
+    if not isinstance(X, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray type, got {type(X)}.")
+
+    if type(axis) not in (tuple, list):
+        axis = (axis,)
+
+    out_ndim = len(axis) + X.ndim
+    axis = normalize_axis_tuple(axis, out_ndim)
+
+    shape_it = iter(X.shape)
+    shape = tuple(1 if ax in axis else next(shape_it) for ax in range(out_ndim))
+
+    return dpt_ext.reshape(X, shape)
+
+
 def repeat(x, repeats, /, *, axis=None):
     """repeat(x, repeats, axis=None)
 

dpnp/dpnp_iface_manipulation.py

@@ -1849,7 +1849,7 @@ def expand_dims(a, axis):
     """
 
     usm_a = dpnp.get_usm_ndarray(a)
-    usm_res = dpt.expand_dims(usm_a, axis=axis)
+    usm_res = dpt_ext.expand_dims(usm_a, axis=axis)
     return dpnp_array._create_from_usm_ndarray(usm_res)