Make normalize_chunks a cubed function

Author: tomwhite

cubed/array/overlap.py

@@ -3,8 +3,7 @@
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.ops import map_selection
 from cubed.types import T_RectangularChunks
-from cubed.utils import _cumsum
-from cubed.vendor.dask.array.core import normalize_chunks
+from cubed.utils import _cumsum, normalize_chunks
 from cubed.vendor.dask.array.overlap import coerce_boundary, coerce_depth
 from cubed.vendor.dask.utils import has_keyword
 

cubed/array_api/creation_functions.py

@@ -11,8 +11,7 @@
     virtual_in_memory,
     virtual_offsets,
 )
-from cubed.utils import normalize_shape, to_chunksize
-from cubed.vendor.dask.array.core import normalize_chunks
+from cubed.utils import normalize_chunks, normalize_shape, to_chunksize
 
 if TYPE_CHECKING:
     from .array_object import Array

cubed/array_api/manipulation_functions.py

@@ -21,8 +21,14 @@
     map_blocks,
     map_selection,
 )
-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.utils import (
+    block_id_to_offset,
+    get_item,
+    normalize_chunks,
+    offset_to_block_id,
+    to_chunksize,
+)
+from cubed.vendor.dask.array.core import broadcast_chunks
 from cubed.vendor.dask.array.reshape import reshape_rechunk
 from cubed.vendor.dask.array.utils import validate_axis
 

cubed/core/array.py

@@ -9,8 +9,7 @@
 from cubed.runtime.types import Callback, Executor
 from cubed.spec import Spec, spec_from_config
 from cubed.storage.zarr import open_if_lazy_zarr_array
-from cubed.utils import array_memory, itemsize
-from cubed.vendor.dask.array.core import normalize_chunks
+from cubed.utils import array_memory, itemsize, normalize_chunks
 
 from .plan import arrays_to_plan
 

cubed/core/groupby.py

@@ -3,8 +3,7 @@
 from cubed.array_api.manipulation_functions import broadcast_to, expand_dims
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.ops import map_blocks, map_selection, reduction
-from cubed.utils import get_item
-from cubed.vendor.dask.array.core import normalize_chunks
+from cubed.utils import get_item, normalize_chunks
 
 if TYPE_CHECKING:
     from cubed.array_api.array_object import Array

cubed/core/indexing.py

@@ -9,8 +9,7 @@
 from cubed.backend_array_api import backend_array_to_numpy_array
 from cubed.core.array import CoreArray
 from cubed.core.ops import general_blockwise, map_selection, merge_chunks
-from cubed.utils import array_size
-from cubed.vendor.dask.array.core import normalize_chunks
+from cubed.utils import array_size, normalize_chunks
 
 if TYPE_CHECKING:
     from cubed.array_api.array_object import Array

cubed/core/ops.py

@@ -31,11 +31,11 @@
     array_size,
     get_item,
     itemsize,
+    normalize_chunks,
     offset_to_block_id,
     to_chunksize,
 )
 from cubed.utils import numblocks as compute_numblocks
-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
 from cubed.vendor.dask.utils import has_keyword

cubed/primitive/blockwise.py

@@ -25,11 +25,11 @@
     chunk_memory,
     get_item,
     map_nested,
+    normalize_chunks,
     split_into,
     to_chunksize,
 )
 from cubed.utils import numblocks as compute_numblocks
-from cubed.vendor.dask.array.core import normalize_chunks
 from cubed.vendor.dask.blockwise import _get_coord_mapping, _make_dims, lol_product
 from cubed.vendor.dask.core import flatten
 

cubed/random.py

@@ -5,8 +5,7 @@
 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, normalize_shape
-from cubed.vendor.dask.array.core import normalize_chunks
+from cubed.utils import block_id_to_offset, normalize_chunks, normalize_shape
 
 
 def random(size, *, dtype=nxp.float64, chunks=None, spec=None):

cubed/utils.py

@@ -13,16 +13,17 @@
 from operator import add, mul
 from pathlib import Path
 from posixpath import join
-from typing import Dict, Tuple, Union, cast
+from typing import Dict, Optional, Tuple, Union, cast
 from urllib.parse import quote, unquote, urlsplit, urlunsplit
 
 import numpy as np
 import tlz as toolz
 from toolz import reduce
 
 from cubed.backend_array_api import namespace as nxp
-from cubed.types import T_DType, T_RectangularChunks, T_RegularChunks, T_Shape
-from cubed.vendor.dask.array.core import _check_regular_chunks, normalize_chunks
+from cubed.types import T_Chunks, T_DType, T_RectangularChunks, T_RegularChunks, T_Shape
+from cubed.vendor.dask.array.core import _check_regular_chunks
+from cubed.vendor.dask.array.core import normalize_chunks as dask_normalize_chunks
 
 PathType = Union[str, Path]
 
@@ -369,3 +370,15 @@ def normalize_dtype(dtype, device=None) -> T_DType:
 
 def itemsize(dtype: T_DType) -> int:
     return dtype.itemsize
+
+
+def normalize_chunks(
+    chunks: T_Chunks,
+    shape: Optional[T_Shape] = None,
+    limit: Optional[int] = None,
+    dtype: Optional[T_DType] = None,
+    previous_chunks: Optional[T_RectangularChunks] = None,
+) -> T_RectangularChunks:
+    return dask_normalize_chunks(
+        chunks, shape=shape, limit=limit, dtype=dtype, previous_chunks=previous_chunks
+    )