ENH: adding casting option to numpy.stack.

numpy/core/shape_base.py

@@ -345,7 +345,7 @@ def hstack(tup):
         return _nx.concatenate(arrs, 1)
 
 
-def _stack_dispatcher(arrays, axis=None, out=None):
+def _stack_dispatcher(arrays, axis=None, out=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 +355,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, casting='same_kind'):
     """
     Join a sequence of arrays along a new axis.
 
@@ -378,6 +378,10 @@ def stack(arrays, axis=0, out=None):
         correct, matching that of what stack would have returned if no
         out argument were specified.
 
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        Controls what kind of data casting may occur. Defaults to 'same_kind'.
+
+
     Returns
     -------
     stacked : ndarray
@@ -430,15 +434,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,
+                           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 +545,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,63 +2,77 @@ 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])
 
 __all__: list[str]
 
+
 @overload
 def atleast_1d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
 @overload
 def atleast_1d(arys: ArrayLike, /) -> NDArray[Any]: ...
 @overload
 def atleast_1d(*arys: ArrayLike) -> list[NDArray[Any]]: ...
 
+
 @overload
 def atleast_2d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
 @overload
 def atleast_2d(arys: ArrayLike, /) -> NDArray[Any]: ...
 @overload
 def atleast_2d(*arys: ArrayLike) -> list[NDArray[Any]]: ...
 
+
 @overload
 def atleast_3d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
 @overload
 def atleast_3d(arys: ArrayLike, /) -> NDArray[Any]: ...
 @overload
 def atleast_3d(*arys: ArrayLike) -> list[NDArray[Any]]: ...
 
+
 @overload
 def vstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
 @overload
 def vstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
 
+
 @overload
 def hstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
 @overload
 def hstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
 
+
 @overload
 def stack(
     arrays: Sequence[_ArrayLike[_SCT]],
     axis: SupportsIndex = ...,
     out: None = ...,
+    casting: None | _CastingKind = ...
 ) -> NDArray[_SCT]: ...
+
+
 @overload
 def stack(
     arrays: Sequence[ArrayLike],
     axis: SupportsIndex = ...,
     out: None = ...,
+    casting: None | _CastingKind = ...
 ) -> NDArray[Any]: ...
+
+
 @overload
 def stack(
     arrays: Sequence[ArrayLike],
     axis: SupportsIndex = ...,
     out: _ArrayType = ...,
+    casting: None | _CastingKind = ...
 ) -> _ArrayType: ...
 
+
 @overload
 def block(arrays: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
 @overload