ENH: adding casting option to numpy.stack.

.gitignore

@@ -173,6 +173,7 @@ numpy/core/src/umath/struct_ufunc_test.c
 numpy/core/src/umath/test_rational.c
 numpy/core/src/umath/umath_tests.c
 numpy/core/src/umath/loops_utils.h
+numpy/core/src/umath/loops_modulo.dispatch.c
 numpy/distutils/__config__.py
 numpy/linalg/umath_linalg.c
 doc/source/**/generated/

doc/release/upcoming_changes/21627.new_feature.rst

@@ -0,0 +1,4 @@
+``casting and dtype`` options for `numpy.stack`
+----------------------------------------------------
+The ``casting and dtype`` options is now available for `numpy.stack`.
+To use it, write ``np.stack(..., dtype=None, casting='same_kind')``.

numpy/core/shape_base.py

@@ -345,7 +345,8 @@ def hstack(tup):
         return _nx.concatenate(arrs, 1)
 
 
-def _stack_dispatcher(arrays, axis=None, out=None):
+def _stack_dispatcher(arrays, axis=None, out=None, *,
+                      dtype=None, casting=None):
     arrays = _arrays_for_stack_dispatcher(arrays, stacklevel=6)
     if out is not None:
         # optimize for the typical case where only arrays is provided
@@ -355,7 +356,7 @@ def _stack_dispatcher(arrays, axis=None, out=None):
 
 
 @array_function_dispatch(_stack_dispatcher)
-def stack(arrays, axis=0, out=None):
+def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"):
     """
     Join a sequence of arrays along a new axis.
 
@@ -378,6 +379,18 @@ def stack(arrays, axis=0, out=None):
         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`.
+
+        .. versionadded:: 1.24
+
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        Controls what kind of data casting may occur. Defaults to 'same_kind'.
+
+        .. versionadded:: 1.24
+
+
     Returns
     -------
     stacked : ndarray
@@ -430,15 +443,17 @@ def stack(arrays, axis=0, out=None):
 
     sl = (slice(None),) * axis + (_nx.newaxis,)
     expanded_arrays = [arr[sl] for arr in arrays]
-    return _nx.concatenate(expanded_arrays, axis=axis, out=out)
+    return _nx.concatenate(expanded_arrays, axis=axis, out=out,
+                           dtype=dtype, casting=casting)
 
 
 # Internal functions to eliminate the overhead of repeated dispatch in one of
 # the two possible paths inside np.block.
 # Use getattr to protect against __array_function__ being disabled.
 _size = getattr(_from_nx.size, '__wrapped__', _from_nx.size)
 _ndim = getattr(_from_nx.ndim, '__wrapped__', _from_nx.ndim)
-_concatenate = getattr(_from_nx.concatenate, '__wrapped__', _from_nx.concatenate)
+_concatenate = getattr(_from_nx.concatenate,
+                       '__wrapped__', _from_nx.concatenate)
 
 
 def _block_format_index(index):
@@ -539,7 +554,7 @@ def _concatenate_shapes(shapes, axis):
     """Given array shapes, return the resulting shape and slices prefixes.
 
     These help in nested concatenation.
-    
+
     Returns
     -------
     shape: tuple of int

numpy/core/shape_base.pyi

@@ -2,7 +2,7 @@ from collections.abc import Sequence
 from typing import TypeVar, overload, Any, SupportsIndex
 
 from numpy import generic
-from numpy._typing import ArrayLike, NDArray, _ArrayLike
+from numpy._typing import ArrayLike, NDArray, _ArrayLike, _CastingKind
 
 _SCT = TypeVar("_SCT", bound=generic)
 _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
@@ -45,18 +45,27 @@ def stack(
     arrays: Sequence[_ArrayLike[_SCT]],
     axis: SupportsIndex = ...,
     out: None = ...,
+    *,
+    dtype: None = ...,
+    casting: None | _CastingKind = ...
 ) -> NDArray[_SCT]: ...
 @overload
 def stack(
     arrays: Sequence[ArrayLike],
     axis: SupportsIndex = ...,
     out: None = ...,
+    *,
+    dtype: None = ...,
+    casting: None | _CastingKind = ...
 ) -> NDArray[Any]: ...
 @overload
 def stack(
     arrays: Sequence[ArrayLike],
     axis: SupportsIndex = ...,
     out: _ArrayType = ...,
+    *,
+    dtype: None = ...,
+    casting: None | _CastingKind = ...
 ) -> _ArrayType: ...
 
 @overload

numpy/core/tests/test_shape_base.py

@@ -449,6 +449,41 @@ def test_stack():
     with assert_warns(FutureWarning):
         result = stack((x for x in range(3)))
     assert_array_equal(result, np.array([0, 1, 2]))
+    #casting and dtype test
+    a = np.array([1, 2, 3])
+    b = np.array([2.5, 3.5, 4.5])
+    res = np.stack((a, b), axis=1, casting="unsafe", dtype=np.int64)
+    expected_res = np.array([[1, 2], [2, 3], [3, 4]])
+    assert_array_equal(res, expected_res)
+    #casting and dtype with TypeError
+    with assert_raises(TypeError):
+        stack((a, b), dtype=np.int64, axis=1, casting="safe")
+
+
+@pytest.mark.parametrize("axis", [0])
+@pytest.mark.parametrize("out_dtype", ["c8", "f4", "f8", ">f8", "i8"])
+@pytest.mark.parametrize("casting",
+                         ['no', 'equiv', 'safe', 'same_kind', 'unsafe'])
+def test_stack_out_and_dtype(axis, out_dtype, casting):
+    to_concat = (array([1, 2]), array([3, 4]))
+    res = array([[1, 2], [3, 4]])
+    out = np.zeros_like(res)
+
+    if not np.can_cast(to_concat[0], out_dtype, casting=casting):
+        with assert_raises(TypeError):
+            stack(to_concat, dtype=out_dtype,
+                  axis=axis, casting=casting)
+    else:
+        res_out = stack(to_concat, out=out,
+                        axis=axis, casting=casting)
+        res_dtype = stack(to_concat, dtype=out_dtype,
+                          axis=axis, casting=casting)
+        assert res_out is out
+        assert_array_equal(out, res_dtype)
+        assert res_dtype.dtype == out_dtype
+
+    with assert_raises(TypeError):
+        stack(to_concat, out=out, dtype=out_dtype, axis=axis)
 
 
 class TestBlock: