Repository: cubed-dev/cubed
Pull request #472 (2024-06-04 14:38:10)

Improve test coverage


Patch:
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index ae7deab9..848587a7 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -10,7 +10,7 @@ from warnings import warn
 
 import numpy as np
 import zarr
-from tlz import concat, partition
+from tlz import concat, first, partition
 from toolz import accumulate, map
 from zarr.indexing import (
     IntDimIndexer,
@@ -37,7 +37,7 @@ from cubed.utils import (
     offset_to_block_id,
     to_chunksize,
 )
-from cubed.vendor.dask.array.core import common_blockdim, normalize_chunks
+from cubed.vendor.dask.array.core import normalize_chunks
 from cubed.vendor.dask.array.utils import validate_axis
 from cubed.vendor.dask.blockwise import broadcast_dimensions, lol_product
 from cubed.vendor.dask.utils import has_keyword
@@ -571,9 +571,6 @@ def map_blocks(
 ) -> "Array":
     """Apply a function to corresponding blocks from multiple input arrays."""
 
-    if drop_axis is None:
-        drop_axis = []
-
     # Handle the case where an array is created by calling `map_blocks` with no input arrays
     if len(args) == 0:
         from cubed.array_api.creation_functions import empty_virtual_array
@@ -1383,7 +1380,9 @@ def unify_chunks(*args: "Array", **kwargs):
         else:
             nameinds.append((a, ind))
 
-    chunkss = broadcast_dimensions(nameinds, blockdim_dict, consolidate=common_blockdim)
+    chunkss = broadcast_dimensions(
+        nameinds, blockdim_dict, consolidate=smallest_blockdim
+    )
 
     arrays = []
     for a, i in arginds:
@@ -1400,8 +1399,36 @@ def unify_chunks(*args: "Array", **kwargs):
             )
             if chunks != a.chunks and all(a.chunks):
                 # this will raise if chunks are not regular
+                # but this should never happen with smallest_blockdim
                 chunksize = to_chunksize(chunks)
                 arrays.append(rechunk(a, chunksize))
             else:
                 arrays.append(a)
     return chunkss, arrays
+
+
+def smallest_blockdim(blockdims):
+    """Find the smallest block dimensions from the list of block dimensions
+
+    Unlike Dask's common_blockdim, this returns regular chunks (assuming
+    regular chunks are passed in).
+    """
+    if not any(blockdims):
+        return ()
+    non_trivial_dims = {d for d in blockdims if len(d) > 1}
+    if len(non_trivial_dims) == 1:
+        return first(non_trivial_dims)
+    if len(non_trivial_dims) == 0:
+        return max(blockdims, key=first)
+
+    if len(set(map(sum, non_trivial_dims))) > 1:
+        raise ValueError("Chunks do not add up to same value", blockdims)
+
+    # find dims with the smallest first chunk
+    m = -1
+    out = None
+    for ntd in non_trivial_dims:
+        if m == -1 or ntd[0] < m:
+            m = ntd[0]
+            out = ntd
+    return out
diff --git a/cubed/primitive/blockwise.py b/cubed/primitive/blockwise.py
index 3c6c0271..1cfc83a9 100644
--- a/cubed/primitive/blockwise.py
+++ b/cubed/primitive/blockwise.py
@@ -665,13 +665,11 @@ def make_blockwise_key_function(
     output: str,
     out_indices: Sequence[Union[str, int]],
     *arrind_pairs: Any,
-    numblocks: Optional[Dict[str, Tuple[int, ...]]] = None,
+    numblocks: Dict[str, Tuple[int, ...]],
     new_axes: Optional[Dict[int, int]] = None,
 ) -> Callable[[List[int]], Any]:
     """Make a function that is the equivalent of make_blockwise_graph."""
 
-    if numblocks is None:
-        raise ValueError("Missing required numblocks argument.")
     new_axes = new_axes or {}
     argpairs = list(toolz.partition(2, arrind_pairs))
 
@@ -723,7 +721,7 @@ def make_blockwise_key_function_flattened(
     output: str,
     out_indices: Sequence[Union[str, int]],
     *arrind_pairs: Any,
-    numblocks: Optional[Dict[str, Tuple[int, ...]]] = None,
+    numblocks: Dict[str, Tuple[int, ...]],
     new_axes: Optional[Dict[int, int]] = None,
 ) -> Callable[[List[int]], Any]:
     # TODO: make this a part of make_blockwise_key_function?
diff --git a/cubed/vendor/dask/array/core.py b/cubed/vendor/dask/array/core.py
index 280e3a28..f92c8dec 100644
--- a/cubed/vendor/dask/array/core.py
+++ b/cubed/vendor/dask/array/core.py
@@ -462,71 +462,3 @@ def _check_regular_chunks(chunkset):
         if chunks[-1] > chunks[0]:
             return False
     return True
-
-
-def common_blockdim(blockdims):
-    """Find the common block dimensions from the list of block dimensions
-
-    Currently only implements the simplest possible heuristic: the common
-    block-dimension is the only one that does not span fully span a dimension.
-    This is a conservative choice that allows us to avoid potentially very
-    expensive rechunking.
-
-    Assumes that each element of the input block dimensions has all the same
-    sum (i.e., that they correspond to dimensions of the same size).
-
-    Examples
-    --------
-    >>> common_blockdim([(3,), (2, 1)])
-    (2, 1)
-    >>> common_blockdim([(1, 2), (2, 1)])
-    (1, 1, 1)
-    >>> common_blockdim([(2, 2), (3, 1)])  # doctest: +SKIP
-    Traceback (most recent call last):
-        ...
-    ValueError: Chunks do not align
-    """
-    if not any(blockdims):
-        return ()
-    non_trivial_dims = {d for d in blockdims if len(d) > 1}
-    if len(non_trivial_dims) == 1:
-        return first(non_trivial_dims)
-    if len(non_trivial_dims) == 0:
-        return max(blockdims, key=first)
-
-    if np.isnan(sum(map(sum, blockdims))):
-        raise ValueError(
-            "Arrays' chunk sizes (%s) are unknown.\n\n"
-            "A possible solution:\n"
-            "  x.compute_chunk_sizes()" % blockdims
-        )
-
-    if len(set(map(sum, non_trivial_dims))) > 1:
-        raise ValueError("Chunks do not add up to same value", blockdims)
-
-    # We have multiple non-trivial chunks on this axis
-    # e.g. (5, 2) and (4, 3)
-
-    # We create a single chunk tuple with the same total length
-    # that evenly divides both, e.g. (4, 1, 2)
-
-    # To accomplish this we walk down all chunk tuples together, finding the
-    # smallest element, adding it to the output, and subtracting it from all
-    # other elements and remove the element itself.  We stop once we have
-    # burned through all of the chunk tuples.
-    # For efficiency's sake we reverse the lists so that we can pop off the end
-    rchunks = [list(ntd)[::-1] for ntd in non_trivial_dims]
-    total = sum(first(non_trivial_dims))
-    i = 0
-
-    out = []
-    while i < total:
-        m = min(c[-1] for c in rchunks)
-        out.append(m)
-        for c in rchunks:
-            c[-1] -= m
-            if c[-1] == 0:
-                c.pop()
-        i += m
-
-    return tuple(out)


Test patch:
diff --git a/cubed/tests/runtime/test_local.py b/cubed/tests/runtime/test_local.py
index 7a1226e9..99e3fad1 100644
--- a/cubed/tests/runtime/test_local.py
+++ b/cubed/tests/runtime/test_local.py
@@ -83,11 +83,10 @@ def test_failure(tmp_path, timing_map, n_tasks, retries, use_backups):
 @pytest.mark.parametrize(
     "timing_map, n_tasks, retries",
     [
-        ({0: [60]}, 10, 2),
+        ({0: [10]}, 10, 2),
     ],
 )
 # fmt: on
-@pytest.mark.skip(reason="This passes, but Python will not exit until the slow task is done.")
 def test_stragglers(tmp_path, timing_map, n_tasks, retries):
     outputs = asyncio.run(
         run_test(
diff --git a/cubed/tests/test_array_api.py b/cubed/tests/test_array_api.py
index e54c4843..49294671 100644
--- a/cubed/tests/test_array_api.py
+++ b/cubed/tests/test_array_api.py
@@ -169,7 +169,7 @@ def test_add(spec, any_executor):
     )
 
 
-def test_add_with_broadcast(spec, executor):
+def test_add_different_chunks(spec, executor):
     a = xp.ones((10, 10), chunks=(10, 2), spec=spec)
     b = xp.ones((10, 10), chunks=(2, 10), spec=spec)
     c = xp.add(a, b)
@@ -178,6 +178,13 @@ def test_add_with_broadcast(spec, executor):
     )
 
 
+def test_add_different_chunks_fail(spec, executor):
+    a = xp.ones((10,), chunks=(3,), spec=spec)
+    b = xp.ones((10,), chunks=(5,), spec=spec)
+    c = xp.add(a, b)
+    assert_array_equal(c.compute(executor=executor), np.ones((10,)) + np.ones((10,)))
+
+
 def test_equal(spec):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
diff --git a/cubed/tests/test_ops.py b/cubed/tests/test_ops.py
new file mode 100644
index 00000000..ae51934b
--- /dev/null
+++ b/cubed/tests/test_ops.py
@@ -0,0 +1,108 @@
+import numpy as np
+import pytest
+from numpy.testing import assert_array_equal
+
+import cubed.array_api as xp
+from cubed.core.ops import smallest_blockdim, unify_chunks
+from cubed.tests.utils import TaskCounter
+
+
+def test_smallest_blockdim():
+    assert smallest_blockdim([]) == ()
+    assert smallest_blockdim([(5,), (5,)]) == (5,)
+    assert smallest_blockdim([(5,), (3, 2)]) == (3, 2)
+    assert smallest_blockdim([(5, 5), (3, 3, 3, 1)]) == (3, 3, 3, 1)
+    assert smallest_blockdim([(2, 1), (2, 1)]) == (2, 1)
+    assert smallest_blockdim([(2, 2, 1), (3, 2), (2, 2, 1)]) == (2, 2, 1)
+
+    with pytest.raises(ValueError, match="Chunks do not add up to same value"):
+        smallest_blockdim([(2, 1), (2, 2)])
+
+
+@pytest.mark.parametrize(
+    "chunks_a, chunks_b, expected_chunksize",
+    [
+        ((2,), (4,), (2,)),
+        ((4,), (2,), (2,)),
+        ((6,), (10,), (6,)),
+        ((10,), (10,), (10,)),
+        ((5,), (10,), (5,)),
+        ((3,), (5,), (3,)),
+        ((5,), (3,), (3,)),
+    ],
+)
+def test_unify_chunks_elemwise(chunks_a, chunks_b, expected_chunksize):
+    a = xp.ones((10,), chunks=chunks_a)
+    b = xp.ones((10,), chunks=chunks_b)
+
+    _, arrays = unify_chunks(a, "i", b, "i")
+    for arr in arrays:
+        assert arr.chunksize == expected_chunksize
+
+    c = xp.add(a, b)
+    assert_array_equal(c.compute(), np.ones((10,)) + np.ones((10,)))
+
+
+@pytest.mark.parametrize(
+    "chunks_a, chunks_b, expected_chunksize",
+    [
+        ((2, 2), (4, 4), (2, 2)),
+        ((2, 4), (4, 2), (2, 2)),
+        ((4, 2), (2, 4), (2, 2)),
+        ((3, 5), (5, 3), (3, 3)),
+        ((3, 10), (10, 3), (3, 3)),
+    ],
+)
+def test_unify_chunks_elemwise_2d(chunks_a, chunks_b, expected_chunksize):
+    a = xp.ones((10, 10), chunks=chunks_a)
+    b = xp.ones((10, 10), chunks=chunks_b)
+
+    _, arrays = unify_chunks(a, "ij", b, "ij")
+    for arr in arrays:
+        assert arr.chunksize == expected_chunksize
+
+    c = xp.add(a, b)
+    assert_array_equal(c.compute(), np.ones((10, 10)) + np.ones((10, 10)))
+
+
+@pytest.mark.parametrize(
+    "chunks_a, chunks_b, expected_chunksize",
+    [
+        ((2, 2), (4, 4), (2, 2)),
+        ((2, 4), (2, 4), (2, 2)),
+        ((4, 2), (4, 2), (2, 2)),
+        ((3, 5), (3, 5), (3, 3)),
+        ((3, 10), (3, 10), (3, 3)),
+    ],
+)
+def test_unify_chunks_blockwise_2d(chunks_a, chunks_b, expected_chunksize):
+    a = xp.ones((10, 10), chunks=chunks_a)
+    b = xp.ones((10, 10), chunks=chunks_b)
+
+    _, arrays = unify_chunks(a, "ij", b, "ji")
+    for arr in arrays:
+        assert arr.chunksize == expected_chunksize
+
+    c = xp.matmul(a, b)
+    assert_array_equal(c.compute(), np.matmul(np.ones((10, 10)), np.ones((10, 10))))
+
+
+def test_unify_chunks_broadcast_scalar():
+    a = xp.ones((10,), chunks=(3,))
+    b = a + 1
+    assert_array_equal(b.compute(), np.ones((10,)) + 1)
+
+
+def test_unify_chunks_broadcast_2d():
+    a = xp.ones((10, 10), chunks=(3, 3))
+    b = xp.ones((10,), chunks=(5,))
+    c = xp.add(a, b)
+
+    # the following checks that b is rechunked *before* broadcasting, to avoid materializing the full (broadcasted) array
+    task_counter = TaskCounter()
+    res = c.compute(callbacks=[task_counter])
+    num_created_arrays = 2  # b rechunked, c
+    # 1 task for rechunk of b, 16 for addition operation
+    assert task_counter.value == num_created_arrays + 1 + 16
+
+    assert_array_equal(res, np.ones((10, 10)) + np.ones((10,)))
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index 7fad3f27..e2f77cc1 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -27,7 +27,10 @@ def spec(tmp_path):
     return cubed.Spec(tmp_path, allowed_mem=100000)
 
 
-def test_fusion(spec):
+@pytest.mark.parametrize(
+    "opt_fn", [None, simple_optimize_dag, multiple_inputs_optimize_dag]
+)
+def test_fusion(spec, opt_fn):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.negative(a)
     c = xp.astype(b, np.float32)
@@ -43,12 +46,20 @@ def test_fusion(spec):
     )
     num_arrays = 2  # a, d
     num_created_arrays = 1  # d (a is not created on disk)
-    assert d.plan.num_arrays(optimize_graph=True) == num_arrays
-    assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
-    assert d.plan.total_nbytes_written(optimize_graph=True) == d.nbytes
+    assert (
+        d.plan.num_arrays(optimize_graph=True, optimize_function=opt_fn) == num_arrays
+    )
+    assert (
+        d.plan.num_tasks(optimize_graph=True, optimize_function=opt_fn)
+        == num_created_arrays + 4
+    )
+    assert (
+        d.plan.total_nbytes_written(optimize_graph=True, optimize_function=opt_fn)
+        == d.nbytes
+    )
 
     task_counter = TaskCounter()
-    result = d.compute(callbacks=[task_counter])
+    result = d.compute(optimize_function=opt_fn, callbacks=[task_counter])
     assert task_counter.value == num_created_arrays + 4
 
     assert_array_equal(
@@ -57,7 +68,10 @@ def test_fusion(spec):
     )
 
 
-def test_fusion_transpose(spec):
+@pytest.mark.parametrize(
+    "opt_fn", [None, simple_optimize_dag, multiple_inputs_optimize_dag]
+)
+def test_fusion_transpose(spec, opt_fn):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.negative(a)
     c = xp.astype(b, np.float32)
@@ -66,10 +80,13 @@ def test_fusion_transpose(spec):
     num_created_arrays = 3  # b, c, d
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 12
     num_created_arrays = 1  # d
-    assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
+    assert (
+        d.plan.num_tasks(optimize_graph=True, optimize_function=opt_fn)
+        == num_created_arrays + 4
+    )
 
     task_counter = TaskCounter()
-    result = d.compute(callbacks=[task_counter])
+    result = d.compute(optimize_function=opt_fn, callbacks=[task_counter])
     assert task_counter.value == num_created_arrays + 4
 
     assert_array_equal(
@@ -81,6 +98,7 @@ def test_fusion_transpose(spec):
 def test_fusion_map_direct(spec):
     # test that operations after a map_direct operation (indexing) can be fused
     # with the map_direct operation
+    # this is only true for the (default) multiple_inputs_optimize_dag optimize function
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = a[1:, :]
     c = xp.negative(b)  # should be fused with b
@@ -102,6 +120,7 @@ def test_fusion_map_direct(spec):
 
 def test_no_fusion(spec):
     # b can't be fused with c because d also depends on b
+    # this is only true for the simple_optimize_dag optimize function
     a = xp.ones((2, 2), chunks=(2, 2), spec=spec)
     b = xp.positive(a)
     c = xp.positive(b)
@@ -126,7 +145,7 @@ def test_no_fusion_multiple_edges(spec):
     c = xp.asarray(b)
     # b and c are the same array, so d has a single dependency
     # with multiple edges
-    # this should not be fused under the current logic
+    # this should not be fused under the current logic in simple_optimize_dag
     d = xp.equal(b, c)
 
     opt_fn = simple_optimize_dag