Repository: cupy/cupy

Issue:
Fusion does not support out argument
Fusion code raises an error when `out` argument is specified.

```
import cupy
import numpy
xp = cupy

a = xp.ones((2,3))
b = xp.ones((2,3)) * 2
z = xp.zeros((2,3))

@cupy.fuse()
def func(a, b, z):
    xp.add(a, b, out=z)

func(a, b, z)
print(z)
```

Results in:
```
Traceback (most recent call last):
  File "test-out.py", line 13, in <module>
    func(a, b, z)
  File "/repos/cupy/cupy/core/fusion.py", line 602, in __call__
    return self._call(*args, **kwargs)
  File "/repos/cupy/cupy/core/fusion.py", line 628, in _call
    self.post_map, self.identity, types)
  File "/repos/cupy/cupy/core/fusion.py", line 499, in _get_fusion
    out_refs = func(*in_refs)
  File "test-out.py", line 11, in func
    xp.add(a, b, out=z)
  File "/repos/cupy/cupy/core/fusion.py", line 710, in __call__
    return _convert(self._fusion_op)(*args, **kwargs)
  File "/repos/cupy/cupy/core/fusion.py", line 339, in res
    var_list.append(_normalize_arg.pop('out'))
AttributeError: 'function' object has no attribute 'pop'
```


Discussion:

Resolving patch:
diff --git a/cupy/core/fusion.py b/cupy/core/fusion.py
--- a/cupy/core/fusion.py
+++ b/cupy/core/fusion.py
@@ -336,10 +336,10 @@ def res(*args, **kwargs):
         mem = get_mem(args)
         var_list = [_normalize_arg(_, mem) for _ in args]
         if 'out' in kwargs:
-            var_list.append(_normalize_arg.pop('out'))
+            var_list.append(_normalize_arg(kwargs.pop('out'), mem))
         if kwargs:
             raise TypeError('Wrong arguments %s' % kwargs)
-        assert nin <= len(var_list) and len(var_list) <= nin + nout
+        assert nin <= len(var_list) <= nin + nout
         in_vars = var_list[:nin]
         out_vars = var_list[nin:]
         can_cast = can_cast1 if _should_use_min_scalar(in_vars) else can_cast2
@@ -360,9 +360,12 @@ def res(*args, **kwargs):
                         v = out_vars[i]
                         ret.append(_FusionRef(v, mem))
                     else:
-                        raise TypeError("Cannot cast from %s to %s"
-                                        % (ty_outs[i], out_vars[i].ty)
-                                        + " with casting rule 'same_kind'")
+                        raise TypeError(
+                            'output (typecode \'{}\') could not be coerced '
+                            'to provided output parameter (typecode \'{}\') '
+                            'according to the casting rule '
+                            '"same_kind"'.format(
+                                ty_outs[i].char, out_vars[i].ty.char))
                 mem.set_op(ufunc.name, op, param_names, nin, nout,
                            in_vars, out_vars, ty_ins + ty_outs)
                 return ret[0] if len(ret) == 1 else tuple(ret)
@@ -498,7 +501,7 @@ def _get_fusion(func, nin, reduce, post_map, identity, input_types, name=None):
     in_refs = [_FusionRef(_, mem) for _ in in_vars]
     out_refs = func(*in_refs)
     out_refs = list(out_refs) if type(out_refs) == tuple else [out_refs]
-    out_refs = filter(lambda i: i is not None, out_refs)
+    out_refs = [_ for _ in out_refs if _ is not None]
     out_refs = [_FusionRef(_normalize_arg(_, mem), mem) for _ in out_refs]
     out_vars = [_normalize_arg(copy(_), mem) for _ in out_refs]
     nout = len(out_vars)


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_fusion.py b/tests/cupy_tests/core_tests/test_fusion.py
--- a/tests/cupy_tests/core_tests/test_fusion.py
+++ b/tests/cupy_tests/core_tests/test_fusion.py
@@ -550,32 +550,87 @@ def random_real(self, lower=-1000, higher=1000):
                                      dtype=numpy.float64,
                                      scale=(higher - lower)) + lower
 
-    def check(self, func, n, gen, *args):
+    def check(self, func, n, gen, args=None):
+        if args is None:
+            args = ((),) * n
         self._check(func, n, gen, args, True)
         self._check(func, n, gen, args, False)
 
-    def _check(self, func, n, gen, args, omit_nin):
+    def _check(self, func, n, gen, args, omit_nin, error_types=None):
+        assert n == len(args), (n, args)
         nin = n if not omit_nin else None
+        if error_types is None:
+            error_types = ()
 
         @cupy.fuse(input_num=nin)
         def f(*x):
             return func(*x)
 
-        if type(gen) == tuple:
-            data = [g(*a) for i, g, a in zip(range(n), list(gen), args)]
+        # Prepare input arrays
+        if not isinstance(gen, tuple):
+            gen = (gen,) * n
+        data0 = tuple([g(*a) for g, a in zip(gen, args)])
+        data1 = tuple([_.copy() for _ in data0])
+
+        # Invoke non-fused function
+        try:
+            ret0 = func(*data0)  # Non-fused
+            err0 = None
+        except Exception as e:
+            if type(e) not in error_types:
+                raise
+            ret0 = None
+            err0 = e
+
+        # Invoke fused function
+        try:
+            ret1 = f(*data1)  # Fused
+            err1 = None
+        except Exception as e:
+            if type(e) not in error_types:
+                raise
+            ret1 = None
+            err1 = e
+
+        self.assertEqual(err0 is None, err1 is None)
+        if err0 is not None:
+            # Error
+            self.assertEqual(type(err0), type(err1))
+            self.assertEqual(str(err0), str(err1))
+            arrs0 = None
+            arrs1 = None
         else:
-            data = [gen(*args) for i in range(n)]
-
-        ret0 = func(*data)  # Non-fused
-        ret1 = f(*data)  # Fused
-        if not isinstance(ret0, tuple):
-            ret0 = [ret0]
-        if not isinstance(ret1, tuple):
-            ret1 = [ret1]
-        for nf, f in zip(ret0, ret1):
-            numpy.testing.assert_array_almost_equal(nf.get(), f.get())
-
-    def check_reduce(self, func, n, reduce_f, gen, *args):
+            # Success
+            self.assertEqual(ret0 is None, ret1 is None)
+            if ret0 is None:
+                # Both return values are None
+                ret0 = ()
+                ret1 = ()
+            else:
+                # Return values must have the same type
+                self.assertEqual(type(ret0), type(ret1))
+
+                if not isinstance(ret0, tuple):
+                    # Single arrays are returned
+                    ret0 = (ret0,)
+                    ret1 = (ret1,)
+                else:
+                    # Tuples are returned
+                    self.assertEqual(len(ret0), len(ret1))
+
+            # Concatenate return values and input arrays
+            arrs0 = ret0 + data0
+            arrs1 = ret1 + data1
+
+            # Test they have same values
+            for nf, f in zip(arrs0, arrs1):
+                numpy.testing.assert_array_almost_equal(nf.get(), f.get())
+
+        return err0 is not None, (arrs0, arrs1)
+
+    def check_reduce(self, func, n, reduce_f, gen, args=None):
+        if args is None:
+            args = ((),) * n
         self._check_reduce(func, n, reduce_f, gen, args, True)
         self._check_reduce(func, n, reduce_f, gen, args, False)
 
@@ -586,18 +641,60 @@ def _check_reduce(self, func, n, reduce_f, gen, args, omit_nin):
         def f(*x):
             return func(*x)
 
-        data = [gen(*args) for i in range(n)]
+        data = [gen(*a) for a in args]
         ret0 = reduce_f(func(*data))  # Non-fused
         ret1 = f(*data)  # Fused
         numpy.testing.assert_array_almost_equal(ret0.get(), ret1.get())
 
+    @testing.for_all_dtypes_combination(
+        names=['src_dtype', 'dst_dtype'], full=True)
+    def test_out_arg(self, src_dtype, dst_dtype):
+
+        def func(x, y, z):
+            return cupy.add(x, y, out=z)
+
+        dtypes = (src_dtype, src_dtype, dst_dtype)
+        ret = self._check(
+            func, 3,
+            lambda iarg: cupy.arange(6).astype(dtypes[iarg]).reshape((2, 3)),
+            [(_,) for _ in range(3)],
+            True, error_types=(TypeError,))
+        is_err, (arrs_n, arrs_f) = ret
+
+        if not is_err:
+            # The returned array must equal to z
+            arr_ret = arrs_f[0]
+            arr_z = arrs_f[3]
+            testing.assert_array_equal(arr_ret, arr_z)
+
+    def test_out_arg2(self):
+
+        def func(x, y, z, u, v):
+            cupy.add(x, y, out=z)
+            cupy.subtract(z, x, out=u)
+            cupy.multiply(z, x, out=v)
+            return u
+
+        ret = self._check(
+            func, 5, self.random_int, ((),) * 5,
+            True, error_types=(TypeError,))
+        is_err, (arrs_n, arrs_f) = ret
+
+        # Must succeed
+        self.assertFalse(is_err)
+
+        # The returned array must equal to u
+        arr_ret = arrs_f[0]
+        arr_u = arrs_f[4]
+        testing.assert_array_equal(arr_ret, arr_u)
+
     def test_bitwise(self):
         self.check(cupy.bitwise_and, 2, self.random_int)
         self.check(cupy.bitwise_or, 2, self.random_int)
         self.check(cupy.bitwise_xor, 2, self.random_int)
         self.check(cupy.invert, 1, self.random_int)
-        self.check(cupy.left_shift, 2, self.random_int, 0, 20)
-        self.check(cupy.right_shift, 2, self.random_int, 0, 20)
+        self.check(cupy.left_shift, 2, self.random_int, ((0, 20),) * 2)
+        self.check(cupy.right_shift, 2, self.random_int, ((0, 20),) * 2)
 
     def test_compare(self):
         self.check(cupy.greater, 2, self.random_int)
@@ -613,17 +710,17 @@ def test_logic_content(self):
         self.check(cupy.isnan, 1, self.random_real)
 
     def test_logic_ops(self):
-        self.check(cupy.logical_and, 2, self.random_int, 0, 2)
-        self.check(cupy.logical_or, 2, self.random_int, 0, 2)
-        self.check(cupy.logical_not, 1, self.random_int, 0, 2)
-        self.check(cupy.logical_xor, 2, self.random_int, 0, 2)
+        self.check(cupy.logical_and, 2, self.random_int, ((0, 2),) * 2)
+        self.check(cupy.logical_or, 2, self.random_int, ((0, 2),) * 2)
+        self.check(cupy.logical_not, 1, self.random_int, ((0, 2),))
+        self.check(cupy.logical_xor, 2, self.random_int, ((0, 2),) * 2)
 
     def test_trigonometric(self):
         self.check(cupy.sin, 1, self.random_real)
         self.check(cupy.cos, 1, self.random_real)
         self.check(cupy.tan, 1, self.random_real)
-        self.check(cupy.arcsin, 1, self.random_real, -1, 1)
-        self.check(cupy.arccos, 1, self.random_real, -1, 1)
+        self.check(cupy.arcsin, 1, self.random_real, ((-1, 1),))
+        self.check(cupy.arccos, 1, self.random_real, ((-1, 1),))
         self.check(cupy.arctan, 1, self.random_real)
         self.check(cupy.hypot, 2, self.random_real)
         self.check(cupy.deg2rad, 1, self.random_real)
@@ -632,12 +729,12 @@ def test_trigonometric(self):
         self.check(cupy.radians, 1, self.random_real)
 
     def test_hyperbolic(self):
-        self.check(cupy.sinh, 1, self.random_real, -10, 10)
-        self.check(cupy.cosh, 1, self.random_real, -10, 10)
-        self.check(cupy.tanh, 1, self.random_real, -10, 10)
-        self.check(cupy.arcsinh, 1, self.random_real, -10, 10)
-        self.check(cupy.arccosh, 1, self.random_real, 1, 10)
-        self.check(cupy.arctanh, 1, self.random_real, 0, 1)
+        self.check(cupy.sinh, 1, self.random_real, ((-10, 10),))
+        self.check(cupy.cosh, 1, self.random_real, ((-10, 10),))
+        self.check(cupy.tanh, 1, self.random_real, ((-10, 10),))
+        self.check(cupy.arcsinh, 1, self.random_real, ((-10, 10),))
+        self.check(cupy.arccosh, 1, self.random_real, ((1, 10),))
+        self.check(cupy.arctanh, 1, self.random_real, ((0, 1),))
 
     def test_rounding(self):
         self.check(cupy.rint, 1, self.random_real)
@@ -646,21 +743,21 @@ def test_rounding(self):
         self.check(cupy.trunc, 1, self.random_real)
 
     def test_explog(self):
-        self.check(cupy.exp, 1, self.random_real, -10, 10)
-        self.check(cupy.expm1, 1, self.random_real, -10, 10)
-        self.check(cupy.exp2, 1, self.random_real, -10, 10)
-        self.check(cupy.log, 1, self.random_real, 0, 10)
-        self.check(cupy.log10, 1, self.random_real, 0, 10)
-        self.check(cupy.log2, 1, self.random_real, 0, 10)
-        self.check(cupy.log1p, 1, self.random_real, -1, 10)
-        self.check(cupy.logaddexp, 2, self.random_real, 0, 10)
-        self.check(cupy.logaddexp2, 2, self.random_real, 0, 10)
+        self.check(cupy.exp, 1, self.random_real, ((-10, 10),))
+        self.check(cupy.expm1, 1, self.random_real, ((-10, 10),))
+        self.check(cupy.exp2, 1, self.random_real, ((-10, 10),))
+        self.check(cupy.log, 1, self.random_real, ((0, 10),))
+        self.check(cupy.log10, 1, self.random_real, ((0, 10),))
+        self.check(cupy.log2, 1, self.random_real, ((0, 10),))
+        self.check(cupy.log1p, 1, self.random_real, ((-1, 10),))
+        self.check(cupy.logaddexp, 2, self.random_real, ((0, 10),) * 2)
+        self.check(cupy.logaddexp2, 2, self.random_real, ((0, 10),) * 2)
 
     def test_floating(self):
         self.check(cupy.signbit, 1, self.random_real)
         self.check(cupy.copysign, 2, self.random_real)
-        self.check(cupy.ldexp, 2, self.random_int, 1, 10)
-        self.check(cupy.frexp, 1, self.random_real, 1, 1000)
+        self.check(cupy.ldexp, 2, self.random_int, ((1, 10),) * 2)
+        self.check(cupy.frexp, 1, self.random_real, ((1, 1000),))
         self.check(cupy.nextafter, 2, self.random_real)
 
     def test_arithmetic(self):
@@ -669,17 +766,17 @@ def test_arithmetic(self):
         self.check(cupy.negative, 1, self.random_real)
         self.check(cupy.multiply, 2, self.random_real)
         self.check(cupy.divide, 2, self.random_real)
-        self.check(cupy.power, 2, self.random_real, 0, 10)
+        self.check(cupy.power, 2, self.random_real, ((0, 10),) * 2)
         self.check(cupy.subtract, 2, self.random_real)
-        self.check(cupy.true_divide, 2, self.random_int, 1, 1000)
-        self.check(cupy.floor_divide, 2, self.random_real, 1, 1000)
+        self.check(cupy.true_divide, 2, self.random_int, ((1, 1000),) * 2)
+        self.check(cupy.floor_divide, 2, self.random_real, ((1, 1000),) * 2)
         self.check(cupy.fmod, 2, self.random_real)
-        self.check(cupy.mod, 2, self.random_int, 1, 1000)
+        self.check(cupy.mod, 2, self.random_int, ((1, 1000),) * 2)
         self.check(cupy.modf, 1, self.random_real)
-        self.check(cupy.remainder, 2, self.random_int, 1, 1000)
+        self.check(cupy.remainder, 2, self.random_int, ((1, 1000),) * 2)
 
     def test_misc(self):
-        self.check(cupy.sqrt, 1, self.random_real, 0, 1000)
+        self.check(cupy.sqrt, 1, self.random_real, ((0, 1000),))
         self.check(cupy.square, 1, self.random_real)
         self.check(cupy.absolute, 1, self.random_real)
         self.check(cupy.abs, 1, self.random_real)
@@ -692,15 +789,15 @@ def test_misc(self):
     def test_special(self):
         self.check(cupy.where, 3,
                    (self.random_bool, self.random_int, self.random_int),
-                   (), (0, 100), (0, 100))
+                   ((), (0, 100), (0, 100)))
         self.check(cupy.clip, 3,
                    (self.random_real, self.random_real, self.random_real),
-                   (0, 1000), (0, 500), (500, 1000))
+                   ((0, 1000), (0, 500), (500, 1000)))
 
     def test_reduce(self):
         self.check_reduce(cupy.bitwise_and, 2, cupy.sum, self.random_int)
-        self.check_reduce(cupy.sqrt, 1, cupy.prod, self.random_int, 1, 2)
-        self.check_reduce(cupy.sqrt, 1, cupy.prod, self.random_real, 1, 2)
+        self.check_reduce(cupy.sqrt, 1, cupy.prod, self.random_int, ((1, 2),))
+        self.check_reduce(cupy.sqrt, 1, cupy.prod, self.random_real, ((1, 2),))
 
         self.check_reduce(lambda x: x, 1, cupy.amax, self.random_int)
         self.check_reduce(lambda x: x, 1, cupy.amin, self.random_int)