Factor out itemsize function

Author: tomwhite

cubed/array_api/array_object.py

@@ -22,7 +22,7 @@
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.array import CoreArray
 from cubed.core.ops import elemwise
-from cubed.utils import memory_repr
+from cubed.utils import itemsize, memory_repr
 from cubed.vendor.dask.widgets import get_template
 
 ARRAY_SVG_SIZE = (
@@ -55,7 +55,7 @@ def _repr_html_(self):
 
         if not math.isnan(self.nbytes):
             nbytes = memory_repr(self.nbytes)
-            cbytes = memory_repr(math.prod(self.chunksize) * self.dtype.itemsize)
+            cbytes = memory_repr(math.prod(self.chunksize) * itemsize(self.dtype))
         else:
             nbytes = "unknown"
             cbytes = "unknown"

cubed/core/array.py

@@ -9,7 +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
+from cubed.utils import array_memory, itemsize
 from cubed.vendor.dask.array.core import normalize_chunks
 
 from .plan import arrays_to_plan
@@ -102,12 +102,12 @@ def size(self):
     @property
     def nbytes(self) -> int:
         """Number of bytes in array"""
-        return self.size * self.dtype.itemsize
+        return self.size * itemsize(self.dtype)
 
     @property
     def itemsize(self) -> int:
         """Length of one array element in bytes"""
-        return self.dtype.itemsize
+        return itemsize(self.dtype)
 
     def _read_stored(self):
         # Only works if the array has been computed

cubed/core/ops.py

@@ -30,6 +30,7 @@
     array_memory,
     array_size,
     get_item,
+    itemsize,
     offset_to_block_id,
     to_chunksize,
 )
@@ -1007,7 +1008,7 @@ def _rechunk_plan(x, chunks, *, min_mem=None):
         shape=x.shape,
         source_chunks=source_chunks,
         target_chunks=target_chunks,
-        itemsize=x.dtype.itemsize,
+        itemsize=itemsize(x.dtype),
         min_mem=min_mem,
         max_mem=rechunker_max_mem,
     )

cubed/core/plan.py

@@ -21,6 +21,7 @@
     chunk_memory,
     extract_stack_summaries,
     is_local_path,
+    itemsize,
     join_path,
     memory_repr,
 )
@@ -584,7 +585,7 @@ def create_zarr_array(lazy_zarr_array, *, config=None):
 def create_zarr_arrays(lazy_zarr_arrays, allowed_mem, reserved_mem):
     # projected memory is size of largest dtype size (for a fill value)
     projected_mem = (
-        max([lza.dtype.itemsize for lza in lazy_zarr_arrays], default=0) + reserved_mem
+        max([itemsize(lza.dtype) for lza in lazy_zarr_arrays], default=0) + reserved_mem
     )
     num_tasks = len(lazy_zarr_arrays)
 

cubed/primitive/rechunk.py

@@ -9,6 +9,7 @@
 from cubed.runtime.types import CubedPipeline
 from cubed.storage.zarr import T_ZarrArray, lazy_zarr_array
 from cubed.types import T_RegularChunks, T_Shape, T_Store
+from cubed.utils import itemsize
 from cubed.vendor.rechunker.algorithm import rechunking_plan
 
 sym_counter = 0
@@ -127,13 +128,12 @@ def _setup_array_rechunk(
     shape = source_array.shape
     source_chunks = source_array.chunks
     dtype = source_array.dtype
-    itemsize = dtype.itemsize
 
     read_chunks, int_chunks, write_chunks = rechunking_plan(
         shape,
         source_chunks,
         target_chunks,
-        itemsize,
+        itemsize(dtype),
         max_mem,
     )
 

cubed/storage/types.py

@@ -4,7 +4,7 @@
 from operator import mul
 
 from cubed.types import T_DType, T_RegularChunks, T_Shape
-from cubed.utils import normalize_dtype
+from cubed.utils import itemsize, normalize_dtype
 
 
 class ArrayMetadata:
@@ -26,7 +26,7 @@ def size(self) -> int:
     @property
     def nbytes(self) -> int:
         """Number of bytes in array"""
-        return self.size * self.dtype.itemsize
+        return self.size * itemsize(self.dtype)
 
     @property
     def _cdata_shape(self) -> T_Shape:

cubed/utils.py

@@ -29,7 +29,7 @@
 
 def array_memory(dtype: T_DType, shape: T_Shape) -> int:
     """Calculate the amount of memory in bytes that an array uses."""
-    return normalize_dtype(dtype).itemsize * prod(shape)
+    return itemsize(normalize_dtype(dtype)) * prod(shape)
 
 
 def chunk_memory(arr) -> int:
@@ -365,3 +365,7 @@ def normalize_dtype(dtype, device=None) -> T_DType:
     """
 
     return np.dtype(dtype)
+
+
+def itemsize(dtype: T_DType) -> int:
+    return dtype.itemsize