diff --git a/cubed/core/ops.py b/cubed/core/ops.py index e589c4e2..36fd9dc0 100644 --- a/cubed/core/ops.py +++ b/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 diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py index 4e092e96..a4070237 100644 --- a/cubed/tests/test_core.py +++ b/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)