Repository: tomwhite/cubed
Pull request #354 (2024-01-20 11:26:42)

Implement `reshape_chunks` using `general_blockwise`


Patch:
diff --git a/cubed/array_api/manipulation_functions.py b/cubed/array_api/manipulation_functions.py
index 5e514a63..a6735799 100644
--- a/cubed/array_api/manipulation_functions.py
+++ b/cubed/array_api/manipulation_functions.py
@@ -1,5 +1,4 @@
 from bisect import bisect
-from itertools import product
 from operator import add, mul
 
 import numpy as np
@@ -8,11 +7,10 @@ from toolz import reduce
 
 from cubed.array_api.creation_functions import empty
 from cubed.backend_array_api import namespace as nxp
-from cubed.backend_array_api import numpy_array_to_backend_array
 from cubed.core import squeeze  # noqa: F401
 from cubed.core import blockwise, rechunk, unify_chunks
 from cubed.core.ops import elemwise, general_blockwise, map_blocks, map_direct
-from cubed.utils import get_item, to_chunksize
+from cubed.utils import block_id_to_offset, get_item, offset_to_block_id, to_chunksize
 from cubed.vendor.dask.array.core import broadcast_chunks, normalize_chunks
 from cubed.vendor.dask.array.reshape import reshape_rechunk
 from cubed.vendor.dask.array.utils import validate_axis
@@ -245,33 +243,35 @@ def reshape_chunks(x, shape, chunks):
     if reduce(mul, shape, 1) != x.size:
         raise ValueError("total size of new array must be unchanged")
 
-    inchunks = normalize_chunks(x.chunks, shape=x.shape, dtype=x.dtype)
+    # TODO: check number of chunks is unchanged
+    # inchunks = normalize_chunks(x.chunks, shape=x.shape, dtype=x.dtype)
     outchunks = normalize_chunks(chunks, shape=shape, dtype=x.dtype)
 
-    # TODO: check number of chunks is unchanged
+    # use an empty template (handles smaller end chunks)
+    template = empty(shape, dtype=x.dtype, chunks=chunks, spec=x.spec)
 
-    # memory allocated by reading one chunk from input array
-    extra_projected_mem = x.chunkmem
+    def block_function(out_key):
+        out_coords = out_key[1:]
+        offset = block_id_to_offset(out_coords, template.numblocks)
+        in_coords = offset_to_block_id(offset, x.numblocks)
+        return (
+            (x.name, *in_coords),
+            (template.name, *out_coords),
+        )
 
-    return map_direct(
+    return general_blockwise(
         _reshape_chunk,
+        block_function,
         x,
+        template,
         shape=shape,
         dtype=x.dtype,
         chunks=outchunks,
-        extra_projected_mem=extra_projected_mem,
-        inchunks=inchunks,
-        outchunks=outchunks,
     )
 
 
-def _reshape_chunk(e, x, inchunks=None, outchunks=None, block_id=None):
-    in_keys = list(product(*[range(len(c)) for c in inchunks]))
-    out_keys = list(product(*[range(len(c)) for c in outchunks]))
-    idx = in_keys[out_keys.index(block_id)]
-    out = x.zarray[get_item(x.chunks, idx)]
-    out = numpy_array_to_backend_array(out)
-    return nxp.reshape(out, e.shape)
+def _reshape_chunk(x, template):
+    return nxp.reshape(x, template.shape)
 
 
 def stack(arrays, /, *, axis=0):
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 2bffc2d3..536b9eea 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -24,7 +24,7 @@ from cubed.core.plan import Plan, new_temp_path
 from cubed.primitive.blockwise import blockwise as primitive_blockwise
 from cubed.primitive.blockwise import general_blockwise as primitive_general_blockwise
 from cubed.primitive.rechunk import rechunk as primitive_rechunk
-from cubed.utils import chunk_memory, get_item, to_chunksize
+from cubed.utils import chunk_memory, get_item, offset_to_block_id, to_chunksize
 from cubed.vendor.dask.array.core import common_blockdim, normalize_chunks
 from cubed.vendor.dask.array.utils import validate_axis
 from cubed.vendor.dask.blockwise import broadcast_dimensions
@@ -535,17 +535,15 @@ def map_blocks(
 
         # Create an array of index offsets with the same chunk structure as the args,
         # which we convert to block ids (chunk coordinates) later.
-        a = args[0]
-        offsets = offsets_virtual_array(a.numblocks, a.spec)
+        arg0 = args[0]
+        numblocks = arg0.numblocks
+        offsets = offsets_virtual_array(numblocks, arg0.spec)
         new_args = args + (offsets,)
 
-        def offset_to_block_id(offset):
-            return np.unravel_index(offset, a.numblocks)
-
         def func_with_block_id(func):
             def wrap(*a, **kw):
                 offset = int(a[-1])  # convert from 0-d array
-                block_id = offset_to_block_id(offset)
+                block_id = offset_to_block_id(offset, numblocks)
                 return func(*a[:-1], block_id=block_id, **kw)
 
             return wrap
diff --git a/cubed/random.py b/cubed/random.py
index 7047f349..9852a229 100644
--- a/cubed/random.py
+++ b/cubed/random.py
@@ -1,12 +1,12 @@
 import random as pyrandom
 
-import numpy as np
 from numpy.random import Generator, Philox
 from zarr.util import normalize_shape
 
 from cubed.backend_array_api import namespace as nxp
 from cubed.backend_array_api import numpy_array_to_backend_array
 from cubed.core.ops import map_blocks
+from cubed.utils import block_id_to_offset
 from cubed.vendor.dask.array.core import normalize_chunks
 
 
@@ -34,8 +34,3 @@ def _random(x, numblocks=None, root_seed=None, block_id=None):
     out = rg.random(x.shape)
     out = numpy_array_to_backend_array(out)
     return out
-
-
-def block_id_to_offset(block_id, numblocks):
-    arr = np.empty(numblocks, dtype=np.int8)
-    return sum((np.array(block_id) * arr.strides).tolist())
diff --git a/cubed/utils.py b/cubed/utils.py
index 819247bc..a9d85cfc 100644
--- a/cubed/utils.py
+++ b/cubed/utils.py
@@ -28,6 +28,16 @@ def chunk_memory(dtype: T_DType, chunksize: T_RegularChunks) -> int:
     return np.dtype(dtype).itemsize * prod(chunksize)
 
 
+def offset_to_block_id(offset: int, numblocks: Tuple[int, ...]) -> Tuple[int, ...]:
+    """Convert an index offset to a block ID (chunk coordinates)."""
+    return tuple(int(i) for i in np.unravel_index(offset, numblocks))
+
+
+def block_id_to_offset(block_id: Tuple[int, ...], numblocks: Tuple[int, ...]) -> int:
+    """Convert a block ID (chunk coordinates) to an index offset."""
+    return int(np.ravel_multi_index(block_id, numblocks))
+
+
 def get_item(chunks: T_RectangularChunks, idx: Tuple[int, ...]) -> Tuple[slice, ...]:
     """Convert a chunk index to a tuple of slices."""
     # could use Dask's cached_cumsum here if it improves performance


Test patch:
diff --git a/cubed/tests/test_array_api.py b/cubed/tests/test_array_api.py
index 365ada4b..8c2f5f63 100644
--- a/cubed/tests/test_array_api.py
+++ b/cubed/tests/test_array_api.py
@@ -480,6 +480,19 @@ def test_reshape_chunks(spec, executor):
     )
 
 
+def test_reshape_chunks_with_smaller_end_chunk(spec, executor):
+    a = xp.arange(10, chunks=4, spec=spec)
+    b = reshape_chunks(a, (2, 5), (2, 2))
+
+    assert b.shape == (2, 5)
+    assert b.chunks == ((2,), (2, 2, 1))
+
+    assert_array_equal(
+        b.compute(executor=executor),
+        np.array([[0, 1, 4, 5, 8], [2, 3, 6, 7, 9]]),
+    )
+
+
 def test_squeeze_1d(spec, executor):
     a = xp.asarray([[1, 2, 3]], chunks=(1, 2), spec=spec)
     b = xp.squeeze(a, 0)
diff --git a/cubed/tests/test_utils.py b/cubed/tests/test_utils.py
index 4c1b0d6d..68743cf0 100644
--- a/cubed/tests/test_utils.py
+++ b/cubed/tests/test_utils.py
@@ -1,15 +1,18 @@
 import inspect
+import itertools
 import platform
 
 import numpy as np
 import pytest
 
 from cubed.utils import (
+    block_id_to_offset,
     chunk_memory,
     extract_stack_summaries,
     join_path,
     map_nested,
     memory_repr,
+    offset_to_block_id,
     peak_measured_mem,
     split_into,
     to_chunksize,
@@ -23,6 +26,19 @@ def test_chunk_memory():
     assert chunk_memory(np.int32, (0,)) == 0
 
 
+def test_block_id_to_offset():
+    numblocks = (5, 3)
+    for block_id in itertools.product(*[list(range(n)) for n in numblocks]):
+        offset = block_id_to_offset(block_id, numblocks)
+        assert offset_to_block_id(offset, numblocks) == block_id
+
+    with pytest.raises(ValueError):
+        block_id_to_offset((6, 12), numblocks)
+
+    with pytest.raises(ValueError):
+        offset_to_block_id(100, numblocks)
+
+
 def test_to_chunksize():
     assert to_chunksize(((3, 3, 3, 1),)) == (3,)
     with pytest.raises(ValueError):