Add unstack to all wrapped libraries

Author: asmeurer

array_api_compat/common/_aliases.py

@@ -485,10 +485,16 @@ def isdtype(
         # array_api_strict implementation will be very strict.
         return dtype == kind
 
+# unstack is a new function in the 2023.12 array API standard
+def unstack(x: ndarray, /, xp, *, axis: int = 0) -> Tuple[ndarray, ...]:
+    if x.ndim == 0:
+        raise ValueError("Input array must be at least 1-d.")
+    return tuple(xp.moveaxis(x, axis, 0))
+
 __all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like',
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'astype', 'std', 'var', 'clip', 'permute_dims', 'reshape',
            'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc', 'matmul',
-           'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']
+           'matrix_transpose', 'tensordot', 'vecdot', 'isdtype', 'unstack']

array_api_compat/cupy/_aliases.py

@@ -112,11 +112,17 @@ def asarray(
     vecdot = cp.vecdot
 else:
     vecdot = get_xp(cp)(_aliases.vecdot)
+
 if hasattr(cp, 'isdtype'):
     isdtype = cp.isdtype
 else:
     isdtype = get_xp(cp)(_aliases.isdtype)
 
+if hasattr(cp, 'unstack'):
+    unstack = cp.unstack
+else:
+    unstack = get_xp(cp)(_aliases.unstack)
+
 __all__ = _aliases.__all__ + ['asarray', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
                               'atanh', 'bitwise_left_shift', 'bitwise_invert',

array_api_compat/dask/array/_aliases.py

@@ -42,6 +42,7 @@
 import dask.array as da
 
 isdtype = get_xp(np)(_aliases.isdtype)
+unstack = get_xp(da)(_aliases.unstack)
 astype = _aliases.astype
 
 # Common aliases

array_api_compat/numpy/_aliases.py

@@ -117,11 +117,17 @@ def asarray(
     vecdot = np.vecdot
 else:
     vecdot = get_xp(np)(_aliases.vecdot)
+
 if hasattr(np, 'isdtype'):
     isdtype = np.isdtype
 else:
     isdtype = get_xp(np)(_aliases.isdtype)
 
+if hasattr(np, 'unstack'):
+    unstack = np.unstack
+else:
+    unstack = get_xp(np)(_aliases.unstack)
+
 __all__ = _aliases.__all__ + ['asarray', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
                               'atanh', 'bitwise_left_shift', 'bitwise_invert',

array_api_compat/torch/_aliases.py

@@ -4,7 +4,8 @@
 from builtins import all as _builtin_all, any as _builtin_any
 
 from ..common._aliases import (matrix_transpose as _aliases_matrix_transpose,
-                               vecdot as _aliases_vecdot, clip as _aliases_clip)
+                               vecdot as _aliases_vecdot, clip as
+                               _aliases_clip, unstack as _aliases_unstack,)
 from .._internal import get_xp
 
 import torch
@@ -191,6 +192,7 @@ def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keep
     return torch.amin(x, axis, keepdims=keepdims)
 
 clip = get_xp(torch)(_aliases_clip)
+unstack = get_xp(torch)(_aliases_unstack)
 
 # torch.sort also returns a tuple
 # https://github.com/pytorch/pytorch/issues/70921
@@ -709,7 +711,7 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
            'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'copysign',
            'divide', 'equal', 'floor_divide', 'greater', 'greater_equal',
            'hypot', 'less', 'less_equal', 'logaddexp', 'multiply', 'not_equal',
-           'pow', 'remainder', 'subtract', 'max', 'min', 'clip', 'sort',
+           'pow', 'remainder', 'subtract', 'max', 'min', 'clip', 'unstack', 'sort',
            'prod', 'sum', 'any', 'all', 'mean', 'std', 'var', 'concat',
            'squeeze', 'broadcast_to', 'flip', 'roll', 'nonzero', 'where',
            'reshape', 'arange', 'eye', 'linspace', 'full', 'ones', 'zeros',