Repository: cubed-dev/cubed
Pull request #546 (2024-08-08 13:31:41)

Rechunk where dict has missing axes
When rechunking with a dict that doesn't contain all axes, then the chunking should be unchanged for those axes that are missing.

In particular, `a.rechunk({})` should be a no-op.

This is consistent with Dask (https://github.com/dask/dask/issues/11261) and Xarray https://github.com/pydata/xarray/pull/9286)


Patch:
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index e6d1ebc9..a4850596 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -756,6 +756,16 @@ def map_direct(
 
 
 def rechunk(x, chunks, target_store=None):
+    if isinstance(chunks, dict):
+        chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
+        for i in range(x.ndim):
+            if i not in chunks:
+                chunks[i] = x.chunks[i]
+            elif chunks[i] is None:
+                chunks[i] = x.chunks[i]
+    if isinstance(chunks, (tuple, list)):
+        chunks = tuple(lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks))
+
     normalized_chunks = normalize_chunks(chunks, x.shape, dtype=x.dtype)
     if x.chunks == normalized_chunks:
         return x


Test patch:
diff --git a/.github/workflows/array-api-tests.yml b/.github/workflows/array-api-tests.yml
index f41649a4..22ca265f 100644
--- a/.github/workflows/array-api-tests.yml
+++ b/.github/workflows/array-api-tests.yml
@@ -79,15 +79,30 @@ jobs:
           array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__isub__(x1, x2)]
           array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__isub__(x, s)]
 
-          # don't run special cases yet
-          array_api_tests/test_special_cases.py
+          # special cases not implemented
+          array_api_tests/test_special_cases.py::test_empty_arrays[std]
+          array_api_tests/test_special_cases.py::test_empty_arrays[var]
+          array_api_tests/test_special_cases.py::test_nan_propagation[std]
+          array_api_tests/test_special_cases.py::test_nan_propagation[var]
 
-          # don't test signatures or names yet as some are not implemented
-          array_api_tests/test_signatures.py
+          # don't test names yet as some are not implemented
           array_api_tests/test_has_names.py
 
-          # very slow
-          array_api_tests/test_creation_functions.py::test_eye
+          # signatures of items not implemented
+          array_api_tests/test_signatures.py::test_func_signature[std]
+          array_api_tests/test_signatures.py::test_func_signature[var]
+          array_api_tests/test_signatures.py::test_func_signature[unique_all]
+          array_api_tests/test_signatures.py::test_func_signature[unique_counts]
+          array_api_tests/test_signatures.py::test_func_signature[unique_inverse]
+          array_api_tests/test_signatures.py::test_func_signature[unique_values]
+          array_api_tests/test_signatures.py::test_func_signature[nonzero]
+          array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
+          array_api_tests/test_signatures.py::test_func_signature[argsort]
+          array_api_tests/test_signatures.py::test_func_signature[sort]
+          array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
+          array_api_tests/test_signatures.py::test_array_method_signature[__dlpack_device__]
+          array_api_tests/test_signatures.py::test_array_method_signature[__setitem__]
+          array_api_tests/test_signatures.py::test_array_method_signature[to_device]
 
           # edge case failures (https://github.com/cubed-dev/cubed/issues/420)
           array_api_tests/test_linalg.py::test_tensordot
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index f40c7a0f..fdf90ca4 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -246,10 +246,19 @@ def test_multiple_ops(spec, executor):
     )
 
 
-@pytest.mark.parametrize("new_chunks", [(1, 2), {0: 1, 1: 2}])
-def test_rechunk(spec, executor, new_chunks):
+@pytest.mark.parametrize(
+    ("new_chunks", "expected_chunks"),
+    [
+        ((1, 2), ((1, 1, 1), (2, 1))),
+        ({0: 1, 1: 2}, ((1, 1, 1), (2, 1))),
+        ({1: 2}, ((2, 1), (2, 1))),  # dim 0 unchanged
+        ({}, ((2, 1), (1, 1, 1))),  # unchanged
+    ],
+)
+def test_rechunk(spec, executor, new_chunks, expected_chunks):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 1), spec=spec)
     b = a.rechunk(new_chunks)
+    assert b.chunks == expected_chunks
     assert_array_equal(
         b.compute(executor=executor),
         np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),