Coerce args to map_blocks to arrays (#566)

Author: tomwhite

cubed/core/ops.py

@@ -578,6 +578,13 @@ def map_blocks(
 ) -> "Array":
     """Apply a function to corresponding blocks from multiple input arrays."""
 
+    from cubed.array_api.creation_functions import asarray
+
+    # Coerce all args to Cubed arrays
+    specs = [a.spec for a in args if hasattr(a, "spec")]
+    spec0 = specs[0] if len(specs) > 0 else spec
+    args = tuple(asarray(a, spec=spec0) for a in args)
+
     # 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

cubed/tests/test_core.py

@@ -235,6 +235,13 @@ def func(x, y):
     assert_array_equal(c.compute(), np.array([[[12, 13]]]))
 
 
+def test_map_blocks_with_non_cubed_array(spec):
+    a = xp.arange(10, dtype="int64", chunks=(2,), spec=spec)
+    b = np.array([1, 2], dtype="int64")  # numpy array will be coerced to cubed
+    c = cubed.map_blocks(nxp.add, a, b, dtype="int64")
+    assert_array_equal(c.compute(), np.array([1, 3, 3, 5, 5, 7, 7, 9, 9, 11]))
+
+
 def test_multiple_ops(spec, executor):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)