Repository: chainer/chainer

Issue:
`optimizer.update()` fails when link has un-initialized parameter
Testing under chainer==6.0.0 with chainerx built and cupy installed.

Code to reproduce
```
import numpy as np

import chainer
import chainer.links as L


class Model(chainer.Chain):

    def __init__(self):
        super(Model, self).__init__()
        with self.init_scope():
            self.l1 = L.Linear(10, 1)
            # l2 is not initialized
            self.l2 = L.Linear(None, 1)

    def forward(self, x):
        # l2 is not used, and never initialized.
        return self.l1(x)


# device = chainer.get_device(-1)  # When using numpy, below code works.
device = chainer.get_device('native')

model = Model()
model.to_device(device)

optimizer = chainer.optimizers.Adam()
optimizer.setup(model)

x = np.zeros((1, 10), dtype=np.float32)
x = device.send(x)
loss = model(x)
loss.backward()
optimizer.update()
```


```
# Traceback (most recent call last):
#   File "tmp3.py", line 34, in <module>
#     optimizer.update()
#   File "/home/user/.pyenv/versions/anaconda3-5.1.0/lib/python3.6/site-packages/chainer/optimizer.py", line 815, in update
#     param.update()
#   File "/home/user/.pyenv/versions/anaconda3-5.1.0/lib/python3.6/site-packages/chainer/variable.py", line 1877, in update
#     self.update_rule.update(self)
#   File "/home/user/.pyenv/versions/anaconda3-5.1.0/lib/python3.6/site-packages/chainer/optimizer.py", line 223, in update
#     self.update_core(param)
#   File "/home/user/.pyenv/versions/anaconda3-5.1.0/lib/python3.6/site-packages/chainer/optimizer.py", line 240, in update_core
#     self.update_core_chainerx(param)
#   File "/home/user/.pyenv/versions/anaconda3-5.1.0/lib/python3.6/site-packages/chainer/optimizer.py", line 283, in update_core_chainerx
#     backend_name = param.array.device.backend.name
# AttributeError: 'NoneType' object has no attribute 'device'
```


Discussion:

Resolving patch:
diff --git a/chainer/optimizer.py b/chainer/optimizer.py
--- a/chainer/optimizer.py
+++ b/chainer/optimizer.py
@@ -8,7 +8,6 @@
 import six
 
 import chainer
-from chainer import backend
 from chainer import link as link_module
 from chainer import optimizer_hooks
 from chainer import serializer as serializer_module
@@ -280,50 +279,51 @@ def update_core_chainerx(self, param):
 
         """
         grad_array = param.grad
-        backend_name = param.array.device.backend.name
-        if backend_name not in ('native', 'cuda'):
-            raise RuntimeError(
-                'Default implementation of Optimizer.update_core_chainerx is '
-                'only provided for native or cuda backends (actual: {}). '
-                'Override Optimizer.update_core_chainerx() to implement '
-                'custom update logic.'.format(backend_name))
+
+        device = param.device
+        fallback_device = device.fallback_device
 
         # Convert state arrays to NumPy/CuPy
-        chainerx_state_arrays = {}
-        for state_name, st in self.state.items():
-            st = self.state[state_name]
-            if isinstance(st, chainerx.ndarray):
-                fallback_arr = backend.from_chx(st)
-                self.state[state_name] = fallback_arr
-                chainerx_state_arrays[state_name] = (st, fallback_arr)
+        chainerx_state_arrays = None
+        state = self.state
+        if state is not None:
+            chainerx_state_arrays = {}
+            for state_name, st in state.items():
+                if isinstance(st, chainerx.ndarray):
+                    fallback_arr = fallback_device.send(st)
+                    state[state_name] = fallback_arr
+                    chainerx_state_arrays[state_name] = (st, fallback_arr)
 
         # Create a temporary parameter with memory-shared NumPy/CuPy array
         # If the ChainerX parameter has a cached NumPy/CuPy copy, use the
         # cache and avoid redundant conversion. Else, create the cache here
         # and use it.
         if param._chainerx_fallback_array is None:
-            param._chainerx_fallback_array = backend.from_chx(
-                param.array)
+            param._chainerx_fallback_array = fallback_device.send(param.array)
 
         temp_param = variable.Variable._init_unchecked(
-            param._chainerx_fallback_array, is_chainerx_array=False)
+            param._chainerx_fallback_array,
+            device=fallback_device,
+            is_chainerx_array=False)
 
         if grad_array is not None:
             temp_param._set_grad_without_check(
-                backend.from_chx(grad_array))
+                fallback_device.send(grad_array))
 
         # Update
         self.update_core(temp_param)
 
         # Restore state arrays
-        for state_name, (arr, fallback_arr) in chainerx_state_arrays.items():
-            cur_arr = self.state[state_name]
-            if cur_arr is not fallback_arr:
-                # The optimizer altered the reference of the state, instead of
-                # updating it in-place. We need to convert the new state back
-                # to ChainerX.
-                arr = backend.to_chx(cur_arr)
-            self.state[state_name] = arr
+        if chainerx_state_arrays:
+            for state_name, (arr, fallback_arr) in (
+                    chainerx_state_arrays.items()):
+                cur_arr = state[state_name]
+                if cur_arr is not fallback_arr:
+                    # The optimizer altered the reference of the state, instead
+                    # of updating it in-place. We need to convert the new state
+                    # back to ChainerX.
+                    arr = device.send(cur_arr)
+                state[state_name] = arr
 
     def init_state(self, param):
         """Initializes the state.


Test patch:
diff --git a/tests/chainer_tests/test_optimizer.py b/tests/chainer_tests/test_optimizer.py
--- a/tests/chainer_tests/test_optimizer.py
+++ b/tests/chainer_tests/test_optimizer.py
@@ -499,6 +499,125 @@ def test_update_chainerx_gpu(self):
         self.check_update()
 
 
+@testing.backend.inject_backend_tests(
+    None,
+    [
+        # NumPy
+        {},
+        # CuPy
+        {'use_cuda': True, 'cuda_device': 0},
+        {'use_cuda': True, 'cuda_device': 1},
+        # ChainerX
+        {'use_chainerx': True, 'chainerx_device': 'native:0'},
+        {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
+        {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
+    ])
+@testing.parameterize(*testing.product({
+    'override_pattern': [
+        'generic',  # only update_core() is overridden
+        'cpu_gpu',  # update_core_{cpu,gpu} are overridden
+        'cpu_gpu_chx',  # update_core_{cpu,gpu,chainerx} are overridden
+    ],
+}))
+class TestGradientMethodUpdate(unittest.TestCase):
+    """Ensures UpdateRule's appropriate methods are called, for various
+    override patterns and parameters with various conditions."""
+
+    def create(self, device):
+
+        class MyLink(chainer.Link):
+            def __init__(self):
+                super(MyLink, self).__init__()
+                with self.init_scope():
+                    self.p1 = chainer.Parameter()  # uninitialized
+
+                    self.p2 = chainer.Parameter(  # initialized, with grad
+                        np.array([3, 2], np.float32))
+                    self.p2.grad = np.array([13, 12], np.float32)
+
+                    self.p3 = chainer.Parameter(  # initialized, without grad
+                        np.array([5, 7], np.float32))
+
+        call_record = []
+        override_pattern = self.override_pattern
+
+        class MyUpdateRule(optimizer.UpdateRule):
+            if override_pattern == 'generic':
+                def update_core(self, param):
+                    call_record.append(('update_core', param))
+
+            elif override_pattern == 'cpu_gpu':
+                def update_core_cpu(self, param):
+                    call_record.append(('update_core_cpu', param))
+
+                def update_core_gpu(self, param):
+                    call_record.append(('update_core_gpu', param))
+
+            elif override_pattern == 'cpu_gpu_chx':
+                def update_core_cpu(self, param):
+                    call_record.append(('update_core_cpu', param))
+
+                def update_core_gpu(self, param):
+                    call_record.append(('update_core_gpu', param))
+
+                def update_core_chainerx(self, param):
+                    call_record.append(('update_core_chainerx', param))
+
+            else:
+                assert False, override_pattern
+
+        class MyOptimizer(optimizer.GradientMethod):
+            def create_update_rule(self):
+                return MyUpdateRule()
+
+        optimizer_ = MyOptimizer()
+        target = MyLink()
+        target.to_device(device)
+        optimizer_.setup(target)
+
+        return optimizer_, call_record
+
+    def test_update(self, backend_config):
+        device = backend_config.device
+        override_pattern = self.override_pattern
+        optimizer, call_record = self.create(device)
+
+        optimizer.update()
+
+        self.assertEqual(len(call_record), 3)
+
+        # Detemine the expected method name that was called.
+        if override_pattern == 'generic':
+            method_name = 'update_core'
+        elif override_pattern == 'cpu_gpu':
+            if isinstance(device, backend.ChainerxDevice):
+                xp = device.fallback_device.xp
+            else:
+                xp = device.xp
+
+            if xp is np:
+                method_name = 'update_core_cpu'
+            else:
+                assert xp is cuda.cupy
+                method_name = 'update_core_gpu'
+        elif override_pattern == 'cpu_gpu_chx':
+            if isinstance(device, backend.ChainerxDevice):
+                method_name = 'update_core_chainerx'
+            elif device.xp is np:
+                method_name = 'update_core_cpu'
+            else:
+                assert device.xp is cuda.cupy
+                method_name = 'update_core_gpu'
+        else:
+            assert False, override_pattern
+
+        # Check call record.
+        # TODO(niboshi): Check the param argument as well.
+        self.assertEqual(call_record[0][0], method_name)
+        self.assertEqual(call_record[1][0], method_name)
+        self.assertEqual(call_record[2][0], method_name)
+
+
 @testing.parameterize(*testing.product({
     'shape': [(4, 3, 2)],
     'dtype': [np.float16, np.float32, np.float64],