Use unstack from array-api-compat (#602)

cubed-dev · Oct 29, 2024 · 1f0f81a · 1f0f81a
1 parent cd358c6
commit 1f0f81a
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/cubed/array_api/manipulation_functions.py b/cubed/array_api/manipulation_functions.py
@@ -463,6 +463,5 @@ def key_function(out_key):
 def _unstack_chunk(*arrs, axis=0):
     # unstack each array in arrs and yield all in turn
     for arr in arrs:
-        # TODO: replace with nxp.unstack(arr, axis=axis) when array-api-compat has unstack
-        for a in tuple(nxp.moveaxis(arr, axis, 0)):
+        for a in nxp.unstack(arr, axis=axis):
             yield a