Repository: pydata/sparse
Pull request #883 (2025-07-12 19:52:43)

Adding `unstack` function


## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] 🪄 Feature
- [ ] 🐞 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📚 Documentation
- [x] 🧪 Test
- [ ] 🛠️ Other


## Checklist

- [x] Code follows style guide
- [x] Tests added
- [ ] Documented the changes

***

## Please explain your changes below.

Adding support for the `unstack` function.

codspeed-hq[bot]: <!-- __CODSPEED_PERFORMANCE_REPORT_COMMENT__ -->
<!-- __CODSPEED_INSTRUMENTATION_PERFORMANCE_REPORT_COMMENT__ -->

## [CodSpeed Performance Report](https://codspeed.io/pydata/sparse/branches/prady0t%3Aunstack-function)

### Merging #883 will **improve performances by 88.05%**

<sub>Comparing <code>prady0t:unstack-function</code> (4906538) with <code>main</code> (c8f8b61)</sub>


#### :tada: Hooray! `pytest-codspeed` just leveled up to 4.0.0!

> A heads-up, this is a **breaking change** and it might affect your current performance baseline a bit. But here's the exciting part - it's packed with new, cool features and promises improved result stability :partying_face:!
> _Curious about what's new? Visit our [releases](https://github.com/CodSpeedHQ/pytest-codspeed/releases) page to delve into all the awesome details about this new version._

### Summary

`⚡ 1` improvements  
`✅ 339` untouched benchmarks  



### Benchmarks breakdown

|     | Benchmark | `BASE` | `HEAD` | Change |
| --- | --------- | ----------------------- | ------------------- | ------ |
| ⚡ | `` test_index_slice[side=100-rank=2-format='gcxs'] `` | 3.6 ms | 1.9 ms | +88.05% |


Patch:
diff --git a/ci/Numba-array-api-xfails.txt b/ci/Numba-array-api-xfails.txt
index dd8ec9e..94a55c2 100644
--- a/ci/Numba-array-api-xfails.txt
+++ b/ci/Numba-array-api-xfails.txt
@@ -70,14 +70,12 @@ array_api_tests/test_creation_functions.py::test_empty_like
 array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 array_api_tests/test_manipulation_functions.py::test_squeeze
 array_api_tests/test_has_names.py::test_has_names[utility-diff]
-array_api_tests/test_has_names.py::test_has_names[manipulation-unstack]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_sum]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_prod]
 array_api_tests/test_has_names.py::test_has_names[indexing-take_along_axis]
 array_api_tests/test_has_names.py::test_has_names[searching-count_nonzero]
 array_api_tests/test_has_names.py::test_has_names[searching-searchsorted]
 array_api_tests/test_signatures.py::test_func_signature[diff]
-array_api_tests/test_signatures.py::test_func_signature[unstack]
 array_api_tests/test_signatures.py::test_func_signature[take_along_axis]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
@@ -107,6 +105,5 @@ array_api_tests/test_searching_functions.py::test_count_nonzero
 array_api_tests/test_searching_functions.py::test_searchsorted
 array_api_tests/test_signatures.py::test_func_signature[cumulative_sum]
 array_api_tests/test_signatures.py::test_func_signature[cumulative_prod]
-array_api_tests/test_manipulation_functions.py::test_unstack
 array_api_tests/test_signatures.py::test_func_signature[count_nonzero]
 array_api_tests/test_signatures.py::test_func_signature[searchsorted]
diff --git a/sparse/numba_backend/__init__.py b/sparse/numba_backend/__init__.py
index efbfa1c..9226da1 100644
--- a/sparse/numba_backend/__init__.py
+++ b/sparse/numba_backend/__init__.py
@@ -128,6 +128,7 @@ from ._common import (
     sum,
     tensordot,
     tile,
+    unstack,
     var,
     vecdot,
     zeros,
@@ -339,6 +340,7 @@ __all__ = [
     "zeros_like",
     "repeat",
     "tile",
+    "unstack",
 ]
 
 
diff --git a/sparse/numba_backend/_common.py b/sparse/numba_backend/_common.py
index 9f1b051..24b1768 100644
--- a/sparse/numba_backend/_common.py
+++ b/sparse/numba_backend/_common.py
@@ -3186,3 +3186,34 @@ def tile(a, reps):
     a = a.reshape(tuple(np.column_stack(([1] * ndim, shape)).reshape(-1)))
     a = a.broadcast_to(tuple(np.column_stack((reps, shape)).reshape(-1)))
     return a.reshape(tuple(np.multiply(reps, shape)))
+
+
+def unstack(x, axis=0):
+    """
+    Splits an array into a sequence of arrays along the given axis.
+
+    Parameters
+    ----------
+    x : SparseArray
+        Input sparse arrays.
+    axis : int
+        Axis along which the array will be split
+
+    Returns
+    -------
+    out : Tuple[SparseArray,...]
+        Tuple of slices along the given dimension. All the arrays have the same shape.
+    """
+    ndim = x.ndim
+
+    if not (-ndim <= axis < ndim):
+        raise ValueError(f"axis must be in range [-{ndim}, {ndim}), got {axis}")
+
+    if not isinstance(x, SparseArray):
+        raise TypeError("`a` must be a SparseArray.")
+
+    if axis < 0:
+        axis = ndim + axis
+    new_order = (axis,) + tuple(i for i in range(ndim) if i != axis)
+    x = x.transpose(new_order)
+    return (*x,)


Test patch:
diff --git a/sparse/numba_backend/tests/test_coo.py b/sparse/numba_backend/tests/test_coo.py
index 767ba83..5074088 100644
--- a/sparse/numba_backend/tests/test_coo.py
+++ b/sparse/numba_backend/tests/test_coo.py
@@ -1981,3 +1981,42 @@ def test_tile(arr, reps):
     result = sparse.tile(sparse_arr, reps).todense()
 
     np.testing.assert_array_equal(result, expected, err_msg=f"Mismatch for shape={arr.shape}, reps={reps}")
+
+
+@pytest.mark.parametrize("ndim", range(1, 5))
+@pytest.mark.parametrize("shape_range", [3])
+def test_unstack_matches_numpy(ndim, shape_range):
+    rng = np.random.default_rng(42)
+    shape = tuple(rng.integers(2, shape_range + 2) for _ in range(ndim))
+    a = rng.integers(0, 10, size=shape)
+    sparse_a = COO.from_numpy(a)
+
+    for axis in range(-ndim, ndim):
+        sparse_parts = sparse.unstack(sparse_a, axis=axis)
+        np_parts = np.moveaxis(a, axis, 0)
+
+        assert len(sparse_parts) == np_parts.shape[0], f"Wrong number of slices on axis {axis}"
+
+        for i, part in enumerate(sparse_parts):
+            expected = np_parts[i]
+            if isinstance(part, COO):
+                actual = part.todense()
+            elif np.isscalar(part):
+                actual = np.array(part)
+            else:
+                raise TypeError(f"Unexpected type returned from unstack: {type(part)}")
+
+            np.testing.assert_array_equal(actual, expected, err_msg=f"Mismatch at slice {i} on axis {axis}")
+
+
+@pytest.mark.parametrize("axis", [-10, 10, 100, -100])
+def test_unstack_invalid_axis(axis):
+    a = COO.from_numpy(np.arange(6).reshape(2, 3))
+    with pytest.raises(ValueError, match="axis must be in range"):
+        sparse.unstack(a, axis)
+
+
+def test_unstack_invalid_type():
+    a = np.arange(6).reshape(2, 3)  # not a sparse array
+    with pytest.raises(TypeError, match="must be a SparseArray"):
+        sparse.unstack(a, axis=0)
diff --git a/sparse/numba_backend/tests/test_namespace.py b/sparse/numba_backend/tests/test_namespace.py
index 40beed8..40ce4fe 100644
--- a/sparse/numba_backend/tests/test_namespace.py
+++ b/sparse/numba_backend/tests/test_namespace.py
@@ -166,6 +166,7 @@ def test_namespace():
         "uint8",
         "unique_counts",
         "unique_values",
+        "unstack",
         "var",
         "vecdot",
         "where",