Skip to content

Implementation of squeeze function #790

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 16, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion dpctl/tensor/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -39,6 +43,7 @@
"reshape",
"permute_dims",
"expand_dims",
"squeeze",
"from_numpy",
"to_numpy",
"asnumpy",
Expand Down
34 changes: 34 additions & 0 deletions dpctl/tensor/_manipulation_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,3 +68,37 @@ 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."
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At the conclusion of iteration, add the line new_shape = tuple(new_shape).

new_shape = tuple(new_shape)
else:
new_shape = tuple(axis for axis in X_shape if axis != 1)
if new_shape == X.shape:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since new_shape is of type list, and X.shape is of type tuple the comparison is always going to return False:

In [1]: x = (1,2,3)

In [2]: y = list(x)

In [3]: y == x
Out[3]: False

return X
else:
return dpt.reshape(X, new_shape)
99 changes: 99 additions & 0 deletions dpctl/tests/test_usm_ndarray_manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,3 +166,102 @@ def test_expand_dims_incorrect_tuple():
pytest.raises(np.AxisError, dpt.expand_dims, X, (0, 5))

pytest.raises(ValueError, dpt.expand_dims, X, (1, 1))


def test_squeeze_incorrect_type():
X_list = list([1, 2, 3, 4, 5])
X_tuple = tuple(X_list)
Xnp = np.array(X_list)

pytest.raises(TypeError, dpt.permute_dims, X_list, 1)
pytest.raises(TypeError, dpt.permute_dims, X_tuple, 1)
pytest.raises(TypeError, dpt.permute_dims, Xnp, 1)


def test_squeeze_0d():
try:
q = dpctl.SyclQueue()
except dpctl.SyclQueueCreationError:
pytest.skip("Queue could not be created")

Xnp = np.array(1)
X = dpt.asarray(Xnp, sycl_queue=q)
Y = dpt.squeeze(X)
Ynp = Xnp.squeeze()
assert_array_equal(Ynp, dpt.asnumpy(Y))

Y = dpt.squeeze(X, 0)
Ynp = Xnp.squeeze(0)
assert_array_equal(Ynp, dpt.asnumpy(Y))

Y = dpt.squeeze(X, (0))
Ynp = Xnp.squeeze((0))
assert_array_equal(Ynp, dpt.asnumpy(Y))

Y = dpt.squeeze(X, -1)
Ynp = Xnp.squeeze(-1)
assert_array_equal(Ynp, dpt.asnumpy(Y))

pytest.raises(np.AxisError, dpt.squeeze, X, 1)
pytest.raises(np.AxisError, dpt.squeeze, X, -2)
pytest.raises(np.AxisError, dpt.squeeze, X, (1))
pytest.raises(np.AxisError, dpt.squeeze, X, (-2))
pytest.raises(ValueError, dpt.squeeze, X, (0, 0))


@pytest.mark.parametrize(
"shapes",
[
(0),
(1),
(1, 2),
(2, 1),
(1, 1),
(2, 2),
(1, 0),
(0, 1),
(1, 2, 1),
(2, 1, 2),
(2, 2, 2),
(1, 1, 1),
(1, 0, 1),
(0, 1, 0),
],
)
def test_squeeze_without_axes(shapes):
try:
q = dpctl.SyclQueue()
except dpctl.SyclQueueCreationError:
pytest.skip("Queue could not be created")

Xnp = np.empty(shapes)
X = dpt.asarray(Xnp, sycl_queue=q)
Y = dpt.squeeze(X)
Ynp = Xnp.squeeze()
assert_array_equal(Ynp, dpt.asnumpy(Y))


@pytest.mark.parametrize("axes", [0, 2, (0), (2), (0, 2)])
def test_squeeze_axes_arg(axes):
try:
q = dpctl.SyclQueue()
except dpctl.SyclQueueCreationError:
pytest.skip("Queue could not be created")

Xnp = np.array([[[1], [2], [3]]])
X = dpt.asarray(Xnp, sycl_queue=q)
Y = dpt.squeeze(X, axes)
Ynp = Xnp.squeeze(axes)
assert_array_equal(Ynp, dpt.asnumpy(Y))


@pytest.mark.parametrize("axes", [1, -2, (1), (-2), (0, 0), (1, 1)])
def test_squeeze_axes_arg_error(axes):
try:
q = dpctl.SyclQueue()
except dpctl.SyclQueueCreationError:
pytest.skip("Queue could not be created")

Xnp = np.array([[[1], [2], [3]]])
X = dpt.asarray(Xnp, sycl_queue=q)
pytest.raises(ValueError, dpt.squeeze, X, axes)