Repository: chainer/chainer

Issue:
Support ChainerX in Variable.addgrad
Required in MNIST data parallel example (#6857)


Discussion:
https://github.com/chainer/chainer/pull/6979 has now been merged.

Resolving patch:
diff --git a/chainer/variable.py b/chainer/variable.py
--- a/chainer/variable.py
+++ b/chainer/variable.py
@@ -1289,12 +1289,18 @@ def addgrad(self, var):
             var (~chainer.Variable): Source variable.
 
         """
-        # TODO(sonots): Implement for ChainerX
-        if self._has_chainerx_array:
-            raise NotImplementedError()
+        dst_device = self.device
+        is_chainerx = dst_device.xp is chainerx
+
+        if is_chainerx != (var.device.xp is chainerx):
+            raise RuntimeError(
+                'Variable.addgrad does not support addition between '
+                'gradients on non-ChainerX and ChainerX devices.\n'
+                'Adding gradient to: {}\n'
+                'Adding gradient from: {}'.format(
+                    dst_device, var.device))
 
-        assert var.xp is not chainerx
-        if var._grad is None:
+        if var.grad is None:
             return
 
         src = var.grad_var
@@ -1303,14 +1309,10 @@ def addgrad(self, var):
             self.initialize(var.shape)
 
         dst = self.grad_var
-
-        src_dev = cuda.get_device_from_array(src.array)
-        dst_dev = cuda.get_device_from_array(self.array)
-
-        if src_dev.id != dst_dev.id:
-            src = chainer.functions.copy(src, dst_dev.id)
-        self._grad_var = src if dst is None else src + dst
-        self._grad = None
+        src_device = src.device
+        if src_device != dst_device:
+            src = chainer.functions.copy(src, dst_device)
+        self.grad_var = src if dst is None else src + dst
 
     def set_creator(self, gen_func):
         """Notifies the variable that the given function is its creator.


Test patch:
diff --git a/tests/chainer_tests/test_link.py b/tests/chainer_tests/test_link.py
--- a/tests/chainer_tests/test_link.py
+++ b/tests/chainer_tests/test_link.py
@@ -530,7 +530,8 @@ def test_addgrads(self):
             l.x = chainer.Parameter(shape=(2, 3),
                                     initializer=initializers.NaN('d'))
             l.y = chainer.Parameter(shape=2)
-            l.u = chainer.Parameter(shape=(2, 3))
+            l.u = chainer.Parameter(shape=(2, 3),
+                                    initializer=initializers.NaN('d'))
             l.v = chainer.Parameter()
         l.x.grad.fill(1)
         l.y.grad.fill(2)
diff --git a/tests/chainer_tests/test_variable.py b/tests/chainer_tests/test_variable.py
--- a/tests/chainer_tests/test_variable.py
+++ b/tests/chainer_tests/test_variable.py
@@ -104,6 +104,23 @@ def muladd(a, b, c):
     return MulAdd().apply((a, b, c))[0]
 
 
+_nonchainerx_backend_params = [
+    # NumPy
+    {},
+    # CuPy
+    {'use_cuda': True, 'cuda_device': 0},
+    {'use_cuda': True, 'cuda_device': 1},
+]
+
+
+_chainerx_backend_params = [
+    # ChainerX
+    {'use_chainerx': True, 'chainerx_device': 'native:0'},
+    {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
+    {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
+]
+
+
 _backend_params = [
     # NumPy
     {},
@@ -826,26 +843,11 @@ def test_pickle_gpu(self):
         cp.testing.assert_array_equal(x.grad, d.grad)
 
 
-@testing.backend.inject_backend_tests(None, _backend_params)
-@testing.backend.inject_backend_tests(None, _backend_params)
-@testing.backend.inject_backend_tests(None, _backend_params)
-@testing.parameterize(*testing.product(
-    {
-        'clear_src_grad,clear_dst_grad': [
-            [False, False],
-            [True, False],
-            [False, True],
-        ],
-    }))
-class TestVariableAddgrad(unittest.TestCase):
+class VariableAddgradTestBase(object):
 
-    def test_addgrad(
-            self, src_backend_config, dst_backend_config,
-            current_backend_config):
-        # TODO(niboshi): Support ChainerX
-        if (src_backend_config.xp is chainerx
-                or dst_backend_config.xp is chainerx):
-            raise unittest.SkipTest('addgrad does not support ChainerX')
+    def check_addgrad(
+            self, should_succeed,
+            src_backend_config, dst_backend_config, current_backend_config):
         src_device = src_backend_config.device
         dst_device = dst_backend_config.device
 
@@ -871,12 +873,72 @@ def test_addgrad(
             b.cleargrad()
 
         with current_backend_config:
-            b.addgrad(a)
+            if should_succeed:
+                b.addgrad(a)
+            else:
+                with pytest.raises(RuntimeError):
+                    b.addgrad(a)
 
-        np_device = chainer.get_device('@numpy')
-        np.testing.assert_array_equal(np_device.send(b.grad), expect_np)
-        assert backend.get_device_from_array(b.data) == dst_device
-        assert backend.get_device_from_array(b.grad) == dst_device
+        if should_succeed:
+            np_device = chainer.get_device('@numpy')
+            np.testing.assert_array_equal(np_device.send(b.grad), expect_np)
+            assert backend.get_device_from_array(b.data) == dst_device
+            assert backend.get_device_from_array(b.grad) == dst_device
+
+
+addgrad_test_parameterize = testing.parameterize(*testing.product(
+    {
+        'clear_src_grad,clear_dst_grad': [
+            [False, False],
+            [True, False],
+            [False, True],
+        ],
+    }))
+
+
+@testing.backend.inject_backend_tests(None, _backend_params)
+@testing.backend.inject_backend_tests(None, _nonchainerx_backend_params)
+@testing.backend.inject_backend_tests(None, _nonchainerx_backend_params)
+@addgrad_test_parameterize
+class TestVariableAddgradNonChainerx(
+        VariableAddgradTestBase, unittest.TestCase):
+
+    def test_addgrad(
+            self, src_backend_config, dst_backend_config,
+            current_backend_config):
+        self.check_addgrad(
+            True, src_backend_config, dst_backend_config,
+            current_backend_config)
+
+
+@testing.backend.inject_backend_tests(None, _backend_params)
+@testing.backend.inject_backend_tests(None, _chainerx_backend_params)
+@testing.backend.inject_backend_tests(None, _chainerx_backend_params)
+@addgrad_test_parameterize
+class TestVariableAddgradChainerx(
+        VariableAddgradTestBase, unittest.TestCase):
+
+    def test_addgrad(
+            self, src_backend_config, dst_backend_config,
+            current_backend_config):
+        self.check_addgrad(
+            True, src_backend_config, dst_backend_config,
+            current_backend_config)
+
+
+@testing.backend.inject_backend_tests(None, _backend_params)
+@testing.backend.inject_backend_tests(None, _nonchainerx_backend_params)
+@testing.backend.inject_backend_tests(None, _chainerx_backend_params)
+@addgrad_test_parameterize
+class TestVariableAddgradBetweenNonChainerxAndChainerx(
+        VariableAddgradTestBase, unittest.TestCase):
+
+    def test_unsupported_addgrad(
+            self, src_backend_config, dst_backend_config,
+            current_backend_config):
+        self.check_addgrad(
+            False, src_backend_config, dst_backend_config,
+            current_backend_config)
 
 
 @testing.parameterize(