Cast inputs to Cubed arrays in apply_ufunc (#551)

* Cast inputs to Cubed arrays in `apply_ufunc`

* Add comment about specs being the same

Author: tomwhite

cubed/core/gufunc.py

@@ -68,7 +68,13 @@ def apply_gufunc(
     # Main code:
 
     # Cast all input arrays to cubed
-    # args = [asarray(a) for a in args]  # TODO: do we need to support casting?
+    # Use a spec if there is one. Note that all args have to have the same spec, and
+    # this will be checked later when constructing the plan (see check_array_specs).
+    from cubed.array_api.creation_functions import asarray
+
+    specs = [a.spec for a in args if hasattr(a, "spec")]
+    spec = specs[0] if len(specs) > 0 else None
+    args = [asarray(a, spec=spec) for a in args]
 
     if len(input_coredimss) != len(args):
         raise ValueError(

cubed/tests/test_gufunc.py

@@ -38,6 +38,16 @@ def add(x, y):
     assert_equal(z, np.array([2, 4, 6]))
 
 
+def test_apply_gufunc_elemwise_01_non_cubed_input(spec):
+    def add(x, y):
+        return x + y
+
+    a = cubed.from_array(np.array([1, 2, 3]), chunks=3, spec=spec)
+    b = np.array([1, 2, 3])
+    z = apply_gufunc(add, "(),()->()", a, b, output_dtypes=a.dtype)
+    assert_equal(z, np.array([2, 4, 6]))
+
+
 def test_apply_gufunc_elemwise_loop(spec):
     def foo(x):
         assert x.shape in ((2,), (1,))