Implemented column_stack, dstack and row_stack

Author: antonwolfy

dpnp/dpnp_iface_manipulation.py

@@ -57,8 +57,10 @@
     "broadcast_arrays",
     "broadcast_to",
     "can_cast",
+    "column_stack",
     "concatenate",
     "copyto",
+    "dstack",
     "expand_dims",
     "flip",
     "fliplr",
@@ -71,6 +73,7 @@
     "result_type",
     "roll",
     "rollaxis",
+    "row_stack",
     "shape",
     "squeeze",
     "stack",
@@ -444,6 +447,61 @@ def can_cast(from_, to, casting="safe"):
     return dpt.can_cast(dtype_from, to, casting)
 
 
+def column_stack(tup):
+    """
+    Stacks 1-D and 2-D arrays as columns into a 2-D array.
+
+    Take a sequence of 1-D arrays and stack them as columns to make a single
+    2-D array. 2-D arrays are stacked as-is, just like with :obj:`dpnp.hstack`.
+    1-D arrays are turned into 2-D columns first.
+
+    For full documentation refer to :obj:`numpy.column_stack`.
+
+    Parameters
+    ----------
+    tup : {dpnp.ndarray, usm_ndarray}
+        A sequence of 1-D or 2-D arrays to stack. All of them must have
+        the same first dimension.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        The array formed by stacking the given arrays.
+
+    See Also
+    --------
+    :obj:`dpnp.stack` : Stack a sequence of arrays along a new axis.
+    :obj:`dpnp.hstack` : Stack arrays in sequence horizontally (column wise).
+    :obj:`dpnp.vstack` : Stack arrays in sequence vertically (row wise).
+    :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.array((1, 2, 3))
+    >>> b = np.array((2, 3, 4))
+    >>> np.column_stack((a, b))
+    array([[1, 2],
+           [2, 3],
+           [3, 4]])
+
+    """
+
+    arrays = []
+    for v in tup:
+        dpnp.check_supported_arrays_type(v)
+
+        if v.ndim == 1:
+            v = v[:, dpnp.newaxis]
+        elif v.ndim != 2:
+            raise ValueError(
+                "Only 1 or 2 dimensional arrays can be column stacked"
+            )
+
+        arrays.append(v)
+    return dpnp.concatenate(arrays, 1)
+
+
 def concatenate(
     arrays, /, *, axis=0, out=None, dtype=None, casting="same_kind"
 ):
@@ -601,6 +659,63 @@ def copyto(dst, src, casting="same_kind", where=True):
         dst_usm[mask_usm] = src_usm[mask_usm]
 
 
+def dstack(tup):
+    """
+    Stack arrays in sequence depth wise (along third axis).
+
+    This is equivalent to concatenation along the third axis after 2-D arrays
+    of shape `(M, N)` have been reshaped to `(M, N, 1)` and 1-D arrays of shape
+    `(N,)` have been reshaped to `(1, N, 1)`. Rebuilds arrays divided by
+    `dsplit`.
+
+    For full documentation refer to :obj:`numpy.dstack`.
+
+    Parameters
+    ----------
+    tup : {dpnp.ndarray, usm_ndarray}
+        One or more array-like sequences. The arrays must have the same shape
+        along all but the third axis. 1-D or 2-D arrays must have the same shape.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        The array formed by stacking the given arrays, will be at least 3-D.
+
+    See Also
+    --------
+    :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis.
+    :obj:`dpnp.stack` : Join a sequence of arrays along a new axis.
+    :obj:`dpnp.block` : Assemble an nd-array from nested lists of blocks.
+    :obj:`dpnp.vstack` : Stack arrays in sequence vertically (row wise).
+    :obj:`dpnp.hstack` : Stack arrays in sequence horizontally (column wise).
+    :obj:`dpnp.column_stack` : Stack 1-D arrays as columns into a 2-D array.
+    :obj:`dpnp.dsplit` : Split array along third axis.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.array((1, 2, 3))
+    >>> b = np.array((2, 3, 4))
+    >>> np.dstack((a, b))
+    array([[[1, 2],
+            [2, 3],
+            [3, 4]]])
+
+    >>> a = np.array([[1], [2], [3]])
+    >>> b = np.array([[2], [3], [4]])
+    >>> np.dstack((a, b))
+    array([[[1, 2]],
+           [[2, 3]],
+           [[3, 4]]])
+
+    """
+
+    arrs = atleast_3d(*tup)
+    if not isinstance(arrs, list):
+        arrs = [arrs]
+    return dpnp.concatenate(arrs, axis=2)
+
+
 def expand_dims(a, axis):
     """
     Expand the shape of an array.
@@ -1738,3 +1853,6 @@ def vstack(tup, *, dtype=None, casting="same_kind"):
     if not isinstance(arrs, list):
         arrs = [arrs]
     return dpnp.concatenate(arrs, axis=0, dtype=dtype, casting=casting)
+
+
+row_stack = vstack

tests/third_party/cupy/manipulation_tests/test_join.py

@@ -1,14 +1,12 @@
-import unittest
-
 import numpy
 import pytest
 
 import dpnp as cupy
+from tests.helper import has_support_aspect64
 from tests.third_party.cupy import testing
 
 
-class TestJoin(unittest.TestCase):
-    @pytest.mark.skip("dpnp.column_stack() is not implemented yet")
+class TestJoin:
     @testing.for_all_dtypes(name="dtype1")
     @testing.for_all_dtypes(name="dtype2")
     @testing.numpy_cupy_array_equal()
@@ -18,21 +16,18 @@ def test_column_stack(self, xp, dtype1, dtype2):
         c = testing.shaped_arange((4, 2), xp, dtype1)
         return xp.column_stack((a, b, c))
 
-    @pytest.mark.skip("dpnp.column_stack() is not implemented yet")
     def test_column_stack_wrong_ndim1(self):
         a = cupy.zeros(())
         b = cupy.zeros((3,))
         with pytest.raises(ValueError):
             cupy.column_stack((a, b))
 
-    @pytest.mark.skip("dpnp.column_stack() is not implemented yet")
     def test_column_stack_wrong_ndim2(self):
         a = cupy.zeros((3, 2, 3))
         b = cupy.zeros((3, 2))
         with pytest.raises(ValueError):
             cupy.column_stack((a, b))
 
-    @pytest.mark.skip("dpnp.column_stack() is not implemented yet")
     def test_column_stack_wrong_shape(self):
         a = cupy.zeros((3, 2))
         b = cupy.zeros((4, 3))
@@ -87,9 +82,8 @@ def test_concatenate_large_4(self, xp, dtype):
         b = testing.shaped_reverse_arange((2, 3, 4), xp, dtype)
         return xp.concatenate((a, b) * 10, axis=-1)
 
-    @pytest.mark.skip("TODO: remove once dpctl #1325 is resolved")
     @testing.for_all_dtypes(name="dtype")
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_array_equal(type_check=has_support_aspect64())
     def test_concatenate_large_5(self, xp, dtype):
         a = testing.shaped_arange((2, 3, 4), xp, dtype)
         b = testing.shaped_reverse_arange((2, 3, 4), xp, "i")
@@ -113,8 +107,7 @@ def test_concatenate_large_f_contiguous(self, xp, dtype):
         e = testing.shaped_arange((2, 3, 2), xp, dtype)
         return xp.concatenate((a, b, c, d, e) * 2, axis=-1)
 
-    @pytest.mark.skip("TODO: remove once dpctl #1325 is resolved")
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_array_equal(type_check=has_support_aspect64())
     def test_concatenate_many_multi_dtype(self, xp):
         a = testing.shaped_arange((2, 1), xp, "i")
         b = testing.shaped_arange((2, 1), xp, "f")
@@ -191,9 +184,8 @@ def test_concatenate_out_invalid_dtype(self):
             with pytest.raises(TypeError):
                 xp.concatenate((a, b, c), axis=1, out=out)
 
-    @pytest.mark.skip("TODO: remove once dpctl #1325 is resolved")
     @testing.for_all_dtypes_combination(names=["dtype1", "dtype2"])
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_array_equal(type_check=has_support_aspect64())
     def test_concatenate_different_dtype(self, xp, dtype1, dtype2):
         a = testing.shaped_arange((3, 4), xp, dtype1)
         b = testing.shaped_arange((3, 4), xp, dtype2)
@@ -235,27 +227,23 @@ def test_concatenate_casting(self, xp, dtype1, dtype2, casting):
         b = testing.shaped_arange((3, 4), xp, dtype1)
         return xp.concatenate((a, b), dtype=dtype2, casting=casting)
 
-    @pytest.mark.skip("dpnp.dstack() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_dstack(self, xp):
         a = testing.shaped_arange((1, 3, 2), xp)
         b = testing.shaped_arange((3,), xp)
         c = testing.shaped_arange((1, 3), xp)
         return xp.dstack((a, b, c))
 
-    @pytest.mark.skip("dpnp.dstack() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_dstack_single_element(self, xp):
         a = testing.shaped_arange((1, 2, 3), xp)
         return xp.dstack((a,))
 
-    @pytest.mark.skip("dpnp.dstack() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_dstack_single_element_2(self, xp):
         a = testing.shaped_arange((1, 2), xp)
         return xp.dstack((a,))
 
-    @pytest.mark.skip("dpnp.dstack() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_dstack_single_element_3(self, xp):
         a = testing.shaped_arange((1,), xp)
@@ -473,33 +461,30 @@ def test_stack_dtype(self, xp, dtype1, dtype2):
     def test_stack_casting(self, xp, dtype1, dtype2, casting):
         a = testing.shaped_arange((3, 4), xp, dtype1)
         b = testing.shaped_arange((3, 4), xp, dtype1)
+        # may raise TypeError or ComplexWarning
         return xp.stack((a, b), dtype=dtype2, casting=casting)
 
-    @pytest.mark.skip("dpnp.row_stack() is not implemented yet")
     @testing.for_all_dtypes(name="dtype1")
     @testing.for_all_dtypes(name="dtype2")
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_array_equal(type_check=has_support_aspect64())
     def test_row_stack(self, xp, dtype1, dtype2):
         a = testing.shaped_arange((4, 3), xp, dtype1)
         b = testing.shaped_arange((3,), xp, dtype2)
         c = testing.shaped_arange((2, 3), xp, dtype1)
         return xp.row_stack((a, b, c))
 
-    @pytest.mark.skip("dpnp.row_stack() is not implemented yet")
     def test_row_stack_wrong_ndim1(self):
         a = cupy.zeros(())
         b = cupy.zeros((3,))
         with pytest.raises(ValueError):
             cupy.row_stack((a, b))
 
-    @pytest.mark.skip("dpnp.row_stack() is not implemented yet")
     def test_row_stack_wrong_ndim2(self):
         a = cupy.zeros((3, 2, 3))
         b = cupy.zeros((3, 2))
         with pytest.raises(ValueError):
             cupy.row_stack((a, b))
 
-    @pytest.mark.skip("dpnp.row_stack() is not implemented yet")
     def test_row_stack_wrong_shape(self):
         a = cupy.zeros((3, 2))
         b = cupy.zeros((4, 3))