Merge pull request numpy#26579 from asmeurer/unstack

ENH: Add unstack()

Author: mattip

doc/release/upcoming_changes/26579.new_function.rst

@@ -0,0 +1,6 @@
+New function `numpy.unstack`
+----------------------------
+
+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`.

doc/source/reference/routines.array-manipulation.rst

@@ -88,6 +88,7 @@ Splitting arrays
    dsplit
    hsplit
    vsplit
+   unstack
 
 Tiling arrays
 =============

numpy/__init__.py

@@ -165,8 +165,8 @@
         str_, subtract, sum, swapaxes, take, tan, tanh, tensordot,
         timedelta64, trace, transpose, true_divide, trunc, typecodes, ubyte,
         ufunc, uint, uint16, uint32, uint64, uint8, uintc, uintp, ulong,
-        ulonglong, unsignedinteger, ushort, var, vdot, vecdot, void, vstack,
-        where, zeros, zeros_like
+        ulonglong, unsignedinteger, unstack, ushort, var, vdot, vecdot, void,
+        vstack, where, zeros, zeros_like
     )
 
     # NOTE: It's still under discussion whether these aliases 

numpy/__init__.pyi

@@ -399,6 +399,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

@@ -1,5 +1,5 @@
 __all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
-           'stack', 'vstack']
+           'stack', 'unstack', 'vstack']
 
 import functools
 import itertools
@@ -11,7 +11,6 @@
 from .multiarray import array, asanyarray, normalize_axis_index
 from . import fromnumeric as _from_nx
 
-
 array_function_dispatch = functools.partial(
     overrides.array_function_dispatch, module='numpy')
 
@@ -261,6 +260,7 @@ def vstack(tup, *, dtype=None, casting="same_kind"):
     dstack : Stack arrays in sequence depth wise (along third axis).
     column_stack : Stack 1-D arrays as columns into a 2-D array.
     vsplit : Split an array into multiple sub-arrays vertically (row-wise).
+    unstack : Split an array into a tuple of sub-arrays along an axis.
 
     Examples
     --------
@@ -331,8 +331,9 @@ def hstack(tup, *, dtype=None, casting="same_kind"):
     vstack : Stack arrays in sequence vertically (row wise).
     dstack : Stack arrays in sequence depth wise (along third axis).
     column_stack : Stack 1-D arrays as columns into a 2-D array.
-    hsplit : Split an array into multiple sub-arrays 
+    hsplit : Split an array into multiple sub-arrays
              horizontally (column-wise).
+    unstack : Split an array into a tuple of sub-arrays along an axis.
 
     Examples
     --------
@@ -414,6 +415,7 @@ def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"):
     concatenate : Join a sequence of arrays along an existing axis.
     block : Assemble an nd-array from nested lists of blocks.
     split : Split array into a list of multiple sub-arrays of equal size.
+    unstack : Split an array into a tuple of sub-arrays along an axis.
 
     Examples
     --------
@@ -456,6 +458,76 @@ 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):
+    return (x,)
+
+@array_function_dispatch(_unstack_dispatcher)
+def unstack(x, /, *, axis=0):
+    """
+    Split an array into a sequence of arrays along the given axis.
+
+    The ``axis`` parameter specifies the dimension along which the array will
+    be split. For example, if ``axis=0`` (the default) it will be the first
+    dimension and if ``axis=-1`` it will be the last dimension.
+
+    The result is a tuple of arrays split along ``axis``.
+
+    .. versionadded:: 2.1.0
+
+    Parameters
+    ----------
+    x : ndarray
+        The array to be unstacked.
+    axis : int, optional
+        Axis along which the array will be split. Default: ``0``.
+
+    Returns
+    -------
+    unstacked : tuple of ndarrays
+        The unstacked arrays.
+
+    See Also
+    --------
+    stack : Join a sequence of arrays along a new axis.
+    concatenate : Join a sequence of arrays along an existing axis.
+    block : Assemble an nd-array from nested lists of blocks.
+    split : Split array into a list of multiple sub-arrays of equal size.
+
+    Notes
+    -----
+    ``unstack`` serves as the reverse operation of :py:func:`stack`, i.e.,
+    ``stack(unstack(x, axis=axis), axis=axis) == x``.
+
+    This function is equivalent to ``tuple(np.moveaxis(x, axis, 0))``, since
+    iterating on an array iterates along the first axis.
+
+    Examples
+    --------
+    >>> arr = np.arange(24).reshape((2, 3, 4))
+    >>> np.unstack(arr)
+    (array([[ 0,  1,  2,  3],
+            [ 4,  5,  6,  7],
+            [ 8,  9, 10, 11]]),
+     array([[12, 13, 14, 15],
+            [16, 17, 18, 19],
+            [20, 21, 22, 23]]))
+    >>> np.unstack(arr, axis=1)
+    (array([[ 0,  1,  2,  3],
+            [12, 13, 14, 15]]),
+     array([[ 4,  5,  6,  7],
+            [16, 17, 18, 19]]),
+     array([[ 8,  9, 10, 11],
+            [20, 21, 22, 23]]))
+    >>> arr2 = np.stack(np.unstack(arr, axis=1), axis=1)
+    >>> arr2.shape
+    (2, 3, 4)
+    >>> np.all(arr == arr2)
+    np.True_
+
+    """
+    if x.ndim == 0:
+        raise ValueError("Input array must be at least 1-d.")
+    return tuple(_nx.moveaxis(x, axis, 0))
 
 # Internal functions to eliminate the overhead of repeated dispatch in one of
 # the two possible paths inside np.block.
@@ -710,7 +782,7 @@ def block(arrays):
     second-last dimension (-2), and so on until the outermost list is reached.
 
     Blocks can be of any dimension, but will not be broadcasted using
-    the normal rules. Instead, leading axes of size 1 are inserted, 
+    the normal rules. Instead, leading axes of size 1 are inserted,
     to make ``block.ndim`` the same for all blocks. This is primarily useful
     for working with scalars, and means that code like ``np.block([v, 1])``
     is valid, where ``v.ndim == 1``.
@@ -756,6 +828,7 @@ def block(arrays):
     dstack : Stack arrays in sequence depth wise (along third axis).
     column_stack : Stack 1-D arrays as columns into a 2-D array.
     vsplit : Split an array into multiple sub-arrays vertically (row-wise).
+    unstack : Split an array into a tuple of sub-arrays along an axis.
 
     Notes
     -----

numpy/_core/shape_base.pyi

@@ -117,6 +117,21 @@ def stack(
     casting: _CastingKind = ...
 ) -> _ArrayType: ...
 
+@overload
+def unstack(
+    array: _ArrayLike[_SCT],
+    /,
+    *,
+    axis: int = ...,
+) -> tuple[NDArray[_SCT], ...]: ...
+@overload
+def unstack(
+    array: ArrayLike,
+    /,
+    *,
+    axis: int = ...,
+) -> tuple[NDArray[Any], ...]: ...
+
 @overload
 def block(arrays: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
 @overload

numpy/_core/tests/test_shape_base.py

@@ -490,6 +490,39 @@ def test_stack():
         stack((a, b), dtype=np.int64, axis=1, casting="safe")
 
 
+def test_unstack():
+    a = np.arange(24).reshape((2, 3, 4))
+
+    for stacks in [np.unstack(a),
+                   np.unstack(a, axis=0),
+                   np.unstack(a, axis=-3)]:
+        assert isinstance(stacks, tuple)
+        assert len(stacks) == 2
+        assert_array_equal(stacks[0], a[0])
+        assert_array_equal(stacks[1], a[1])
+
+    for stacks in [np.unstack(a, axis=1),
+                   np.unstack(a, axis=-2)]:
+        assert isinstance(stacks, tuple)
+        assert len(stacks) == 3
+        assert_array_equal(stacks[0], a[:, 0])
+        assert_array_equal(stacks[1], a[:, 1])
+        assert_array_equal(stacks[2], a[:, 2])
+
+    for stacks in [np.unstack(a, axis=2),
+                   np.unstack(a, axis=-1)]:
+        assert isinstance(stacks, tuple)
+        assert len(stacks) == 4
+        assert_array_equal(stacks[0], a[:, :, 0])
+        assert_array_equal(stacks[1], a[:, :, 1])
+        assert_array_equal(stacks[2], a[:, :, 2])
+        assert_array_equal(stacks[3], a[:, :, 3])
+
+    assert_raises(ValueError, np.unstack, a, axis=3)
+    assert_raises(ValueError, np.unstack, a, axis=-4)
+    assert_raises(ValueError, np.unstack, np.array(0), axis=0)
+
+
 @pytest.mark.parametrize("axis", [0])
 @pytest.mark.parametrize("out_dtype", ["c8", "f4", "f8", ">f8", "i8"])
 @pytest.mark.parametrize("casting",

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

@@ -53,3 +53,6 @@ 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, axis=0), tuple[npt.NDArray[np.int64], ...])
+assert_type(np.unstack(AR_LIKE_f8, axis=0), tuple[npt.NDArray[Any], ...])