Add squeeze func

Author: vlad-perevezentsev

dpctl/tensor/__init__.py

@@ -25,7 +25,11 @@
 from dpctl.tensor._ctors import asarray, empty
 from dpctl.tensor._device import Device
 from dpctl.tensor._dlpack import from_dlpack
-from dpctl.tensor._manipulation_functions import expand_dims, permute_dims
+from dpctl.tensor._manipulation_functions import (
+    expand_dims,
+    permute_dims,
+    squeeze,
+)
 from dpctl.tensor._reshape import reshape
 from dpctl.tensor._usmarray import usm_ndarray
 
@@ -39,6 +43,7 @@
     "reshape",
     "permute_dims",
     "expand_dims",
+    "squeeze",
     "from_numpy",
     "to_numpy",
     "asnumpy",

dpctl/tensor/_manipulation_functions.py

@@ -84,3 +84,36 @@ def expand_dims(X, axes):
     shape = tuple(1 if ax in axes else next(shape_it) for ax in range(out_ndim))
 
     return dpt.reshape(X, shape)
+
+
+def squeeze(X, axes=None):
+    """
+    squeeze(X: usm_ndarray, axes: int or tuple or list) -> usm_ndarray
+
+    Removes singleton dimensions (axes) from X; returns a view, if possible,
+    a copy otherwise, but with all or a subset of the dimensions
+    of length 1 removed.
+    """
+    if not isinstance(X, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray type, got {type(X)}.")
+    X_shape = X.shape
+    if axes is not None:
+        if not isinstance(axes, (tuple, list)):
+            axes = (axes,)
+        axes = normalize_axis_tuple(axes, X.ndim if X.ndim != 0 else X.ndim + 1)
+        new_shape = []
+        for i, x in enumerate(X_shape):
+            if i not in axes:
+                new_shape.append(x)
+            else:
+                if x != 1:
+                    raise ValueError(
+                        "Cannot select an axis to squeeze out "
+                        "which has size not equal to one."
+                    )
+    else:
+        new_shape = [axis for axis in X_shape if axis != 1]
+    if new_shape == X.shape:
+        return X
+    else:
+        return dpt.reshape(X, tuple(new_shape))