Make itemsize support other Array API implementations (#804)

Author: tomwhite

cubed/backend_array_api.py

@@ -59,5 +59,17 @@ def numpy_array_to_backend_array(arr, *, dtype=None):
     return namespace.asarray(arr, dtype=dtype)
 
 
+def backend_dtype_to_numpy_dtype(dtype):
+    if isinstance(dtype, np.dtype):
+        return dtype
+    elif isinstance(dtype, list):
+        return np.dtype(
+            [(field[0], backend_dtype_to_numpy_dtype(field[1])) for field in dtype]
+        )
+    else:
+        a = namespace.empty((), dtype=dtype)
+        return np.dtype(backend_array_to_numpy_array(a).dtype)
+
+
 # jax doesn't support in-place assignment, so we use .at[].set() instead.
 IS_IMMUTABLE_ARRAY = "jax" in xp_name

cubed/tests/test_utils.py

@@ -14,6 +14,7 @@
     extract_stack_summaries,
     is_cloud_storage_path,
     is_local_path,
+    itemsize,
     join_path,
     map_nested,
     memory_repr,
@@ -32,6 +33,15 @@ def test_array_memory():
     assert array_memory(np.int32, (0,)) == 0
 
 
+def test_itemsize():
+    assert itemsize(np.bool) == 1
+    assert itemsize(np.int32) == 4
+    assert itemsize(np.int64) == 8
+    assert itemsize(np.dtype(np.int64)) == 8
+    assert itemsize([("a", np.int32), ("b", np.int64)]) == 12
+    assert itemsize(np.dtype([("a", np.int32), ("b", np.int64)])) == 12
+
+
 def test_block_id_to_offset():
     numblocks = (5, 3)
     for block_id in itertools.product(*[list(range(n)) for n in numblocks]):

cubed/utils.py

@@ -1,4 +1,5 @@
 import collections
+import inspect
 import itertools
 import numbers
 import platform
@@ -20,6 +21,7 @@
 import tlz as toolz
 from toolz import reduce
 
+from cubed.backend_array_api import backend_dtype_to_numpy_dtype
 from cubed.backend_array_api import namespace as nxp
 from cubed.types import T_Chunks, T_DType, T_RectangularChunks, T_RegularChunks, T_Shape
 from cubed.vendor.dask.array.core import _check_regular_chunks
@@ -369,7 +371,14 @@ def normalize_dtype(dtype, device=None) -> T_DType:
 
 
 def itemsize(dtype: T_DType) -> int:
-    return dtype.itemsize
+    """Return the length of one array element in bytes."""
+    if hasattr(dtype, "itemsize") and not inspect.isdatadescriptor(dtype.itemsize):
+        return dtype.itemsize
+    elif isinstance(dtype, list):
+        return sum(itemsize(v) for _, v in dtype)
+    else:
+        # if dtype has no itemsize property, convert to numpy and use its itemsize
+        return backend_dtype_to_numpy_dtype(dtype).itemsize
 
 
 def normalize_chunks(