From bb20ae1326b93cc32eb98e9dfdf794fea39d111d Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Mon, 14 Oct 2024 12:26:03 +0200 Subject: [PATCH 1/4] Add implementation of dpnp.unstack() --- dpnp/dpnp_iface_manipulation.py | 85 +++++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py index 1e847366d14f..594dd8a45586 100644 --- a/dpnp/dpnp_iface_manipulation.py +++ b/dpnp/dpnp_iface_manipulation.py @@ -95,6 +95,7 @@ "transpose", "trim_zeros", "unique", + "unstack", "vsplit", "vstack", ] @@ -1721,6 +1722,8 @@ def hstack(tup, *, dtype=None, casting="same_kind"): :obj:`dpnp.block` : Assemble an ndarray from nested lists of blocks. :obj:`dpnp.split` : Split array into a list of multiple sub-arrays of equal size. + :obj:`dpnp.unstack` : Split an array into a tuple of sub-arrays along + an axis. Examples -------- @@ -2759,6 +2762,8 @@ def stack(arrays, /, *, axis=0, out=None, dtype=None, casting="same_kind"): :obj:`dpnp.block` : Assemble an ndarray from nested lists of blocks. :obj:`dpnp.split` : Split array into a list of multiple sub-arrays of equal size. + :obj:`dpnp.unstack` : Split an array into a tuple of sub-arrays along + an axis. Examples -------- @@ -3259,6 +3264,84 @@ def unique( return _unpack_tuple(result) +def unstack(x, /, *, axis=0): + """ + Split an array into a sequence of arrays along the given axis. + + The `axis` parameter specifies the dimension along which the array will + be split. For example, if `axis=0` (the default) it will be the first + dimension and if `axis=-1` it will be the last dimension. + + The result is a tuple of arrays split along `axis`. + + For full documentation refer to :obj:`numpy.unstack`. + + Parameters + ---------- + x : {dpnp.ndarray, usm_ndarray} + The array to be unstacked. + axis : int, optional + Axis along which the array will be split. + Default: ``0``. + + Returns + ------- + unstacked : tuple of dpnp.ndarray + The unstacked arrays. + + See Also + -------- + :obj:`dpnp.stack` : Join a sequence of arrays along a new axis. + :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis. + :obj:`dpnp.block` : Assemble an ndarray from nested lists of blocks. + :obj:`dpnp.split` : Split array into a list of multiple sub-arrays of equal + size. + + Notes + ----- + ``unstack`` serves as the reverse operation of :obj:`dpnp.stack`, i.e., + ``dpnp.stack(dpnp.unstack(x, axis=axis), axis=axis) == x``. + + This function is equivalent to ``tuple(dpnp.moveaxis(x, axis, 0))``, since + iterating on an array iterates along the first axis. + + Examples + -------- + >>> import dpnp as np + >>> arr = np.arange(24).reshape((2, 3, 4)) + >>> np.unstack(arr) + (array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11]]), + array([[12, 13, 14, 15], + [16, 17, 18, 19], + [20, 21, 22, 23]])) + + >>> np.unstack(arr, axis=1) + (array([[ 0, 1, 2, 3], + [12, 13, 14, 15]]), + array([[ 4, 5, 6, 7], + [16, 17, 18, 19]]), + array([[ 8, 9, 10, 11], + [20, 21, 22, 23]])) + + >>> arr2 = np.stack(np.unstack(arr, axis=1), axis=1) + >>> arr2.shape + (2, 3, 4) + >>> np.all(arr == arr2) + array(True) + + """ + + usm_x = dpnp.get_usm_ndarray(x) + + if usm_x.ndim == 0: + raise ValueError("Input array must be at least 1-d.") + + res = dpt.unstack(usm_x, axis=axis) + return tuple(dpnp_array._create_from_usm_ndarray(a) for a in res) + + def vsplit(ary, indices_or_sections): """ Split an array into multiple sub-arrays vertically (row-wise). @@ -3367,6 +3450,8 @@ def vstack(tup, *, dtype=None, casting="same_kind"): :obj:`dpnp.block` : Assemble an ndarray from nested lists of blocks. :obj:`dpnp.split` : Split array into a list of multiple sub-arrays of equal size. + :obj:`dpnp.unstack` : Split an array into a tuple of sub-arrays along + an axis. Examples -------- From 2f86a9739289b366932a954d651c9aceacb11dad Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Mon, 14 Oct 2024 12:27:03 +0200 Subject: [PATCH 2/4] Update manipulation.rst --- doc/reference/manipulation.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/reference/manipulation.rst b/doc/reference/manipulation.rst index a9682c488bf4..509b7db847c8 100644 --- a/doc/reference/manipulation.rst +++ b/doc/reference/manipulation.rst @@ -95,6 +95,7 @@ Joining arrays dpnp.dstack dpnp.column_stack dpnp.row_stack + dpnp.unstack Splitting arrays From ef6aea8d8d64e65e92a86ffd6f0c0c03e4010ec1 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Mon, 14 Oct 2024 13:03:32 +0200 Subject: [PATCH 3/4] Add TestUnstack to test_arraymanipulation.py --- tests/test_arraymanipulation.py | 79 +++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/tests/test_arraymanipulation.py b/tests/test_arraymanipulation.py index e8bc95574620..00eebd079dc8 100644 --- a/tests/test_arraymanipulation.py +++ b/tests/test_arraymanipulation.py @@ -822,6 +822,85 @@ def test_generator(self): dpnp.stack(map(lambda x: x, dpnp.ones((3, 2)))) +# numpy.unstack() is available since numpy >= 2.1 +@testing.with_requires("numpy>=2.1") +class TestUnstack: + def test_non_array_input(self): + with pytest.raises(TypeError): + dpnp.unstack(1) + + @pytest.mark.parametrize( + "input", [([1, 2, 3],), [dpnp.int32(1), dpnp.int32(2), dpnp.int32(3)]] + ) + def test_scalar_input(self, input): + with pytest.raises(TypeError): + dpnp.unstack(input) + + @pytest.mark.parametrize("dtype", get_all_dtypes()) + def test_0d_array_input(self, dtype): + np_a = numpy.array(1, dtype=dtype) + dp_a = dpnp.array(np_a, dtype=dtype) + + with pytest.raises(ValueError): + numpy.unstack(np_a) + with pytest.raises(ValueError): + dpnp.unstack(dp_a) + + @pytest.mark.parametrize("dtype", get_all_dtypes()) + def test_1d_array(self, dtype): + np_a = numpy.array([1, 2, 3], dtype=dtype) + dp_a = dpnp.array(np_a, dtype=dtype) + + np_res = numpy.unstack(np_a) + dp_res = dpnp.unstack(dp_a) + assert len(dp_res) == len(np_res) + for dp_arr, np_arr in zip(dp_res, np_res): + assert_array_equal(dp_arr.asnumpy(), np_arr) + + @pytest.mark.parametrize("dtype", get_all_dtypes()) + def test_2d_array(self, dtype): + np_a = numpy.array([[1, 2, 3], [4, 5, 6]], dtype=dtype) + dp_a = dpnp.array(np_a, dtype=dtype) + + np_res = numpy.unstack(np_a, axis=0) + dp_res = dpnp.unstack(dp_a, axis=0) + assert len(dp_res) == len(np_res) + for dp_arr, np_arr in zip(dp_res, np_res): + assert_array_equal(dp_arr.asnumpy(), np_arr) + + @pytest.mark.parametrize("axis", [0, 1, -1]) + @pytest.mark.parametrize("dtype", get_all_dtypes()) + def test_2d_array_axis(self, axis, dtype): + np_a = numpy.array([[1, 2, 3], [4, 5, 6]], dtype=dtype) + dp_a = dpnp.array(np_a, dtype=dtype) + + np_res = numpy.unstack(np_a, axis=axis) + dp_res = dpnp.unstack(dp_a, axis=axis) + assert len(dp_res) == len(np_res) + for dp_arr, np_arr in zip(dp_res, np_res): + assert_array_equal(dp_arr.asnumpy(), np_arr) + + @pytest.mark.parametrize("axis", [2, -3]) + @pytest.mark.parametrize("dtype", get_all_dtypes()) + def test_invalid_axis(self, axis, dtype): + np_a = numpy.array([[1, 2, 3], [4, 5, 6]], dtype=dtype) + dp_a = dpnp.array(np_a, dtype=dtype) + + with pytest.raises(AxisError): + numpy.unstack(np_a, axis=axis) + with pytest.raises(AxisError): + dpnp.unstack(dp_a, axis=axis) + + @pytest.mark.parametrize("dtype", get_all_dtypes()) + def test_empty_array(self, dtype): + np_a = numpy.array([], dtype=dtype) + dp_a = dpnp.array(np_a, dtype=dtype) + + np_res = numpy.unstack(np_a) + dp_res = dpnp.unstack(dp_a) + assert len(dp_res) == len(np_res) + + class TestVstack: def test_non_iterable(self): assert_raises(TypeError, dpnp.vstack, 1) From c2a69d001eb9c5ca80b4973b1aedd6bad7e9c499 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Mon, 14 Oct 2024 16:04:33 +0200 Subject: [PATCH 4/4] Apply reviw remarks --- doc/reference/manipulation.rst | 1 - dpnp/dpnp_iface_manipulation.py | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/doc/reference/manipulation.rst b/doc/reference/manipulation.rst index 1bc22560ac38..2fb9e21bb622 100644 --- a/doc/reference/manipulation.rst +++ b/doc/reference/manipulation.rst @@ -95,7 +95,6 @@ Joining arrays dpnp.dstack dpnp.column_stack dpnp.row_stack - dpnp.unstack Splitting arrays diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py index 5999dd310092..98d64380c950 100644 --- a/dpnp/dpnp_iface_manipulation.py +++ b/dpnp/dpnp_iface_manipulation.py @@ -3370,8 +3370,8 @@ def unstack(x, /, *, axis=0): Split an array into a sequence of arrays along the given axis. The `axis` parameter specifies the dimension along which the array will - be split. For example, if `axis=0` (the default) it will be the first - dimension and if `axis=-1` it will be the last dimension. + be split. For example, if ``axis=0`` (the default) it will be the first + dimension and if ``axis=-1`` it will be the last dimension. The result is a tuple of arrays split along `axis`. @@ -3400,8 +3400,8 @@ def unstack(x, /, *, axis=0): Notes ----- - ``unstack`` serves as the reverse operation of :obj:`dpnp.stack`, i.e., - ``dpnp.stack(dpnp.unstack(x, axis=axis), axis=axis) == x``. + :obj:`dpnp.unstack` serves as the reverse operation of :obj:`dpnp.stack`, + i.e., ``dpnp.stack(dpnp.unstack(x, axis=axis), axis=axis) == x``. This function is equivalent to ``tuple(dpnp.moveaxis(x, axis, 0))``, since iterating on an array iterates along the first axis.