Review changes

Author: mtsokol

doc/release/upcoming_changes/26579.new_function.rst

@@ -1,5 +1,6 @@
-New function ``unstack()``
---------------------------
+New function `numpy.unstack`
+----------------------------
 
-A new function ``unstack()``, which splits an array into a tuple of arrays
-along an axis. It serves as the inverse of ``np.stack()``.
+A new function ``np.unstack(array, axis=...)`` was added, which splits
+an array into a tuple of arrays along an axis. It serves as the inverse
+of `numpy.stack`.

numpy/__init__.pyi

@@ -397,6 +397,7 @@ from numpy._core.shape_base import (
     hstack as hstack,
     stack as stack,
     vstack as vstack,
+    unstack as unstack,
 )
 
 from numpy.lib import (

numpy/_core/shape_base.py

@@ -457,7 +457,7 @@ def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"):
     return _nx.concatenate(expanded_arrays, axis=axis, out=out,
                            dtype=dtype, casting=casting)
 
-def _unstack_dispatcher(x, *, axis=None):
+def _unstack_dispatcher(x, /, *, axis=None):
     return (x,)
 
 @array_function_dispatch(_unstack_dispatcher)

numpy/_core/shape_base.pyi

@@ -120,8 +120,17 @@ def stack(
 @overload
 def unstack(
     array: _ArrayLike[_SCT],
-    axis: int = 0
+    /,
+    *,
+    axis: int = ...,
 ) -> tuple[NDArray[_SCT], ...]: ...
+@overload
+def unstack(
+    array: ArrayLike,
+    /,
+    *,
+    axis: int = ...,
+) -> tuple[NDArray[Any], ...]: ...
 
 @overload
 def block(arrays: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...

numpy/typing/tests/data/reveal/shape_base.pyi

@@ -54,5 +54,5 @@ assert_type(np.kron(AR_f8, AR_f8), npt.NDArray[np.floating[Any]])
 assert_type(np.tile(AR_i8, 5), npt.NDArray[np.int64])
 assert_type(np.tile(AR_LIKE_f8, [2, 2]), npt.NDArray[Any])
 
-assert_type(np.unstack(AR_i8, 5), list[npt.NDArray[np.int64]])
-assert_type(np.unstack(AR_LIKE_f8, 5), list[npt.NDArray[Any]])
+assert_type(np.unstack(AR_i8, axis=0), list[npt.NDArray[np.int64]])
+assert_type(np.unstack(AR_LIKE_f8, axis=0), list[npt.NDArray[Any]])