Repository: chainer/chainer

Issue:
Make `Variable.grad.setter` call `arr.require_grad()` for ChainerX array in order to align the double-backprop behavior with NumPy/CuPy
Original issue: The following code fails at the 2 commented lines. It succeeds for `numpy` and `cupy`.

```py
import chainerx as xp
import chainer
from chainer import functions as F


x_arr = xp.array([1, 2], 'float32')
gy_arr = xp.array([3, 4], 'float32')
ggx_arr = xp.array([5, 6], 'float32')

x = chainer.Variable(x_arr, requires_grad=True)
y = F.identity(x)
y.grad = gy_arr

gy = y.grad_var
y.backward(enable_double_backprop=True)


assert x.grad_var.requires_grad is True  # fail
x.grad_var.grad = ggx_arr  # error
x.grad_var.backward()


print(gy.grad_var)
```



Discussion:
`y.grad_var.requires_grad` should be `True` after `y.grad = gy_arr`.

Resolving patch:
diff --git a/chainer/variable.py b/chainer/variable.py
--- a/chainer/variable.py
+++ b/chainer/variable.py
@@ -912,8 +912,15 @@ def data(self, d):
 
         self.array = d
 
-    def _set_chainerx_grad(self, g):
-        # Assigns chainerx.ndarray.grad
+    def _set_chainerx_grad(self, g, from_grad_var):
+        # Assigns chainerx.ndarray.grad.
+        #
+        # If the main array is connected to the graph, in order to enable
+        # double-backprop, the grad will also be backprop-required
+        # (a view is created not to affect the given grad).
+        # If the given grad is from a grad_var, this operation is skipped,
+        # as the status of the given grad reflects the necessity of
+        # double-backprop.
         assert self.xp is chainerx
         if not self._requires_grad and g is not None:
             raise RuntimeError(
@@ -925,11 +932,17 @@ def _set_chainerx_grad(self, g):
                 raise RuntimeError(
                     'Cannot set a gradient to an empty variable')
         elif arr.is_backprop_required():
+            # If g is grad-stopped, require grad on it.
+            # Make a view in order not to affect the input.
+            if (g is not None
+                    and not from_grad_var
+                    and not g.is_backprop_required()):
+                g = g.view().require_grad()
             arr.set_grad(g)
 
     def _set_grad_without_check(self, g):
         if self._has_chainerx_array:
-            self._set_chainerx_grad(g)
+            self._set_chainerx_grad(g, False)
             self._grad_var = None
             return
 
@@ -986,7 +999,9 @@ def grad(self, g):
 
     def _set_grad_var_without_check(self, gv):
         if self._has_chainerx_array:
-            self._set_chainerx_grad(None if gv is None else gv._data[0])
+            self._set_chainerx_grad(
+                None if gv is None else gv._data[0],
+                True)
             self._grad_var = gv
             return
 
@@ -1189,8 +1204,8 @@ def _to_device(self, device, allow_unchaining):
         else:
             self._data = [new_arr]
             if grad_var is not None:
-                grad_var.to_device(device)
-                # _grad has been invalidated by grad_var.to_device().
+                grad_var._to_device(device, allow_unchaining=allow_unchaining)
+                # _grad has been invalidated by the line above.
                 self._grad = grad_var.array
 
         # ensure that the node tracks the device migration


Test patch:
diff --git a/tests/chainer_tests/functions_tests/array_tests/test_copy.py b/tests/chainer_tests/functions_tests/array_tests/test_copy.py
--- a/tests/chainer_tests/functions_tests/array_tests/test_copy.py
+++ b/tests/chainer_tests/functions_tests/array_tests/test_copy.py
@@ -94,11 +94,6 @@ def test_double_backward(self, src_backend_config, dst_backend_config):
 
         y_var = functions.copy(x_var, dst_device)
 
-        # TODO(niboshi): Remove this workround after Variable.grad.setter is
-        # fixed so that it calls gy.require_grad() internally.
-        if dst_backend_config.xp is chainerx:
-            gy.require_grad()
-
         y_var.grad = gy
 
         gy_var = y_var.grad_var
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
@@ -757,7 +757,12 @@ def test_copydata_from_to_uninitialized_parameters(
 
 
 @testing.backend.inject_backend_tests(None, _backend_params)
-@testing.parameterize(*testing.product({'shape': [(10,), (0,), ()]}))
+@testing.parameterize(*testing.product(
+    {
+        'shape': [(10,), (0,), ()],
+        'requires_grad': [True, False],
+    }
+))
 class TestVariableGrad(unittest.TestCase):
 
     def test_grad(self, backend_config):
@@ -765,22 +770,48 @@ def test_grad(self, backend_config):
             np.random.uniform(-1, 1, self.shape).astype(np.float32))
         g = backend_config.get_array(
             np.random.uniform(0.1, 10, self.shape).astype(np.float32))
-        v = chainer.Variable(x)
-        v.grad = g
-        backend_config.xp.testing.assert_array_equal(v.grad, g)
+        v = chainer.Variable(x, requires_grad=self.requires_grad)
+        expected_error = (
+            backend_config.xp is chainerx
+            and not self.requires_grad)
 
-    def test_grad_var(self, backend_config):
+        if expected_error:
+            with pytest.raises(Exception):
+                v.grad = g
+        else:
+            v.grad = g
+
+            assert v.grad_var.requires_grad is True
+            assert v.grad is not None
+            assert v.requires_grad == self.requires_grad
+            backend_config.xp.testing.assert_array_equal(v.grad, g)
+
+    def check_grad_var(self, backend_config, grad_var_requires_grad):
         x = backend_config.get_array(
             np.random.uniform(-1, 1, self.shape).astype(np.float32))
         g = backend_config.get_array(
             np.random.uniform(0.1, 10, self.shape).astype(np.float32))
-        v = chainer.Variable(x)
-        gv = chainer.Variable(g)
-        v.grad_var = gv
-        backend_config.xp.testing.assert_array_equal(v.grad, g)
+        v = chainer.Variable(x, requires_grad=self.requires_grad)
+        gv = chainer.Variable(g, requires_grad=grad_var_requires_grad)
+        expected_error = (
+            backend_config.xp is chainerx
+            and not self.requires_grad)
 
-        # Same instance should be returned each time.
-        assert v.grad_var is gv
+        if expected_error:
+            with pytest.raises(Exception):
+                v.grad_var = gv
+        else:
+            v.grad_var = gv
+
+            assert v.requires_grad == self.requires_grad
+            backend_config.xp.testing.assert_array_equal(v.grad, g)
+
+            # Same instance should be returned each time.
+            assert v.grad_var is gv
+
+    def test_grad_var(self, backend_config):
+        self.check_grad_var(backend_config, True)
+        self.check_grad_var(backend_config, False)
 
 
 @testing.backend.inject_backend_tests(None, _backend_params)