Added support dtype and out arguments for dpnp.concatenate and dpnp.stack functions (#1608)

* Added support dtype and out args for concatenate and out functions

* address comments

---------

Co-authored-by: Anton <100830759+antonwolfy@users.noreply.github.com>

Author: npolina4

dpnp/dpnp_iface_manipulation.py

@@ -452,19 +452,29 @@ def concatenate(
 
     For full documentation refer to :obj:`numpy.concatenate`.
 
+    Parameters
+    ----------
+    arrays : {dpnp.ndarray, usm_ndarray}
+        The arrays must have the same shape, except in the dimension corresponding
+        to axis (the first, by default).
+    axis : int, optional
+        The axis along which the arrays will be joined. If axis is None, arrays are
+        flattened before use. Default is 0.
+    out : dpnp.ndarray, optional
+        If provided, the destination to place the result. The shape must be correct,
+        matching that of what concatenate would have returned if no out argument were
+        specified.
+    dtype : str or dtype
+        If provided, the destination array will have this dtype. Cannot be provided
+        together with out.
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        Controls what kind of data casting may occur. Defaults to 'same_kind'.
+
     Returns
     -------
     out : dpnp.ndarray
         The concatenated array.
 
-    Limitations
-    -----------
-    Each array in `arrays` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
-    will be raised.
-    Parameters `out` and `dtype` are supported with default value.
-    Otherwise the function will be executed sequentially on CPU.
-
     See Also
     --------
     :obj:`dpnp.array_split` : Split an array into multiple sub-arrays of equal or near-equal size.
@@ -496,25 +506,20 @@ def concatenate(
 
     """
 
-    if out is not None:
-        pass
-    elif dtype is not None:
-        pass
-    elif casting != "same_kind":
-        pass
-    else:
-        usm_arrays = [dpnp.get_usm_ndarray(x) for x in arrays]
-        usm_res = dpt.concat(usm_arrays, axis=axis)
-        return dpnp_array._create_from_usm_ndarray(usm_res)
-
-    return call_origin(
-        numpy.concatenate,
-        arrays,
-        axis=axis,
-        out=out,
-        dtype=dtype,
-        casting=casting,
-    )
+    if dtype is not None and out is not None:
+        raise TypeError(
+            "concatenate() only takes `out` or `dtype` as an argument, but both were provided."
+        )
+
+    usm_arrays = [dpnp.get_usm_ndarray(x) for x in arrays]
+    usm_res = dpt.concat(usm_arrays, axis=axis)
+    res = dpnp_array._create_from_usm_ndarray(usm_res)
+    if dtype is not None:
+        res = res.astype(dtype, casting=casting, copy=False)
+    elif out is not None:
+        dpnp.copyto(out, res, casting=casting)
+        return out
+    return res
 
 
 def copyto(dst, src, casting="same_kind", where=True):
@@ -868,19 +873,21 @@ def hstack(tup, *, dtype=None, casting="same_kind"):
 
     For full documentation refer to :obj:`numpy.hstack`.
 
+    Parameters
+    ----------
+    tup : {dpnp.ndarray, usm_ndarray}
+        The arrays must have the same shape along all but the second axis,
+        except 1-D arrays which can be any length.
+    dtype : str or dtype
+        If provided, the destination array will have this dtype.
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        Controls what kind of data casting may occur. Defaults to 'same_kind'.
+
     Returns
     -------
     out : dpnp.ndarray
         The stacked array which has one more dimension than the input arrays.
 
-    Limitations
-    -----------
-    Each array in `tup` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
-    will be raised.
-    Parameters `dtype` and `casting` are supported with default value.
-    Otherwise the function will be executed sequentially on CPU.
-
     See Also
     --------
     :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis.
@@ -1357,26 +1364,32 @@ def squeeze(a, /, axis=None):
     )
 
 
-def stack(arrays, /, *, axis=0, out=None, dtype=None, **kwargs):
+def stack(arrays, /, *, axis=0, out=None, dtype=None, casting="same_kind"):
     """
     Join a sequence of arrays along a new axis.
 
     For full documentation refer to :obj:`numpy.stack`.
 
+    Parameters
+    ----------
+    arrays : {dpnp.ndarray, usm_ndarray}
+        Each array must have the same shape.
+    axis : int, optional
+        The axis in the result array along which the input arrays are stacked.
+    out : dpnp.ndarray, optional
+        If provided, the destination to place the result. The shape must be correct,
+        matching that of what stack would have returned if no out argument were specified.
+    dtype : str or dtype
+        If provided, the destination array will have this dtype. Cannot be provided
+        together with out.
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        Controls what kind of data casting may occur. Defaults to 'same_kind'.
+
     Returns
     -------
     out : dpnp.ndarray
         The stacked array which has one more dimension than the input arrays.
 
-    Limitations
-    -----------
-    Each array in `arrays` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
-    will be raised.
-    Parameters `out` and `dtype` are supported with default value.
-    Keyword argument `kwargs` is currently unsupported.
-    Otherwise the function will be executed sequentially on CPU.
-
     See Also
     --------
     :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis.
@@ -1409,25 +1422,20 @@ def stack(arrays, /, *, axis=0, out=None, dtype=None, **kwargs):
 
     """
 
-    if kwargs:
-        pass
+    if dtype is not None and out is not None:
+        raise TypeError(
+            "stack() only takes `out` or `dtype` as an argument, but both were provided."
+        )
+
+    usm_arrays = [dpnp.get_usm_ndarray(x) for x in arrays]
+    usm_res = dpt.stack(usm_arrays, axis=axis)
+    res = dpnp_array._create_from_usm_ndarray(usm_res)
+    if dtype is not None:
+        res = res.astype(dtype, casting=casting, copy=False)
     elif out is not None:
-        pass
-    elif dtype is not None:
-        pass
-    else:
-        usm_arrays = [dpnp.get_usm_ndarray(x) for x in arrays]
-        usm_res = dpt.stack(usm_arrays, axis=axis)
-        return dpnp_array._create_from_usm_ndarray(usm_res)
-
-    return call_origin(
-        numpy.stack,
-        arrays,
-        axis=axis,
-        out=out,
-        dtype=dtype,
-        **kwargs,
-    )
+        dpnp.copyto(out, res, casting=casting)
+        return out
+    return res
 
 
 def swapaxes(a, axis1, axis2):
@@ -1649,19 +1657,21 @@ def vstack(tup, *, dtype=None, casting="same_kind"):
 
     For full documentation refer to :obj:`numpy.vstack`.
 
+    Parameters
+    ----------
+    tup : {dpnp.ndarray, usm_ndarray}
+        The arrays must have the same shape along all but the first axis.
+        1-D arrays must have the same length.
+    dtype : str or dtype
+        If provided, the destination array will have this dtype.
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        Controls what kind of data casting may occur. Defaults to 'same_kind'.
+
     Returns
     -------
     out : dpnp.ndarray
         The array formed by stacking the given arrays, will be at least 2-D.
 
-    Limitations
-    -----------
-    Each array in `tup` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
-    will be raised.
-    Parameters `dtype` and `casting` are supported with default value.
-    Otherwise the function will be executed sequentially on CPU.
-
     See Also
     --------
     :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis.

tests/test_arraymanipulation.py

@@ -305,7 +305,6 @@ def test_concatenate_3d(self, dtype):
         dp_res = dpnp.concatenate((dp_a0.T, dp_a1.T, dp_a2.T), axis=0)
         assert_array_equal(dp_res.asnumpy(), np_res)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @pytest.mark.parametrize(
         "dtype", get_all_dtypes(no_bool=True, no_none=True)
     )
@@ -329,7 +328,6 @@ def test_concatenate_out(self, dtype):
         assert_array_equal(dp_out.asnumpy(), np_out)
         assert_array_equal(dp_res.asnumpy(), np_res)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @pytest.mark.parametrize(
         "dtype", get_all_dtypes(no_bool=True, no_none=True)
     )
@@ -487,7 +485,6 @@ def test_empty_arrays_input(self, dtype):
         dp_res = dpnp.stack(dp_arrays, axis=1)
         assert_array_equal(dp_res.asnumpy(), np_res)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @pytest.mark.parametrize("dtype", get_all_dtypes())
     def test_out(self, dtype):
         np_a = numpy.array([1, 2, 3], dtype=dtype)
@@ -536,7 +533,6 @@ def test_generator_input(self):
         with pytest.raises(TypeError):
             dpnp.stack((x for x in range(3)))
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @pytest.mark.usefixtures("suppress_complex_warning")
     @pytest.mark.parametrize("arr_dtype", get_all_dtypes())
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
@@ -552,7 +548,6 @@ def test_casting_dtype(self, arr_dtype, dtype):
         dp_res = dpnp.stack((dp_a, dp_b), axis=1, casting="unsafe", dtype=dtype)
         assert_array_equal(dp_res.asnumpy(), np_res)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @pytest.mark.parametrize("arr_dtype", get_float_complex_dtypes())
     @pytest.mark.parametrize("dtype", [dpnp.bool, dpnp.int32, dpnp.int64])
     def test_invalid_casting_dtype(self, arr_dtype, dtype):
@@ -939,3 +934,17 @@ def test_can_cast():
     assert dpnp.can_cast(X, "float32") == numpy.can_cast(X_np, "float32")
     assert dpnp.can_cast(X, dpnp.int32) == numpy.can_cast(X_np, numpy.int32)
     assert dpnp.can_cast(X, dpnp.int64) == numpy.can_cast(X_np, numpy.int64)
+
+
+def test_concatenate_out_dtype():
+    x = dpnp.ones((5, 5))
+    out = dpnp.empty_like(x)
+    with pytest.raises(TypeError):
+        dpnp.concatenate([x], out=out, dtype="i4")
+
+
+def test_stack_out_dtype():
+    x = dpnp.ones((5, 5))
+    out = dpnp.empty_like(x)
+    with pytest.raises(TypeError):
+        dpnp.stack([x], out=out, dtype="i4")