Repository: chainer/chainer

Issue:
chainer.training.ParallelUpdater does not properly work w/ cupy.core.elementwise
**"Data-parallel Computation on Multiple GPUs with `Trainer` (and `ParallelUpdater`)"**, which is one of topic in [the official document](https://docs.chainer.org/en/stable/guides/gpu.html), does not work properly when the model contains any function-call of `cupy.core.elementwise`.

In _cupy/cupy/core/elementwise.pxi line 127_, it checks if the device between the current device and the device of the array. However, the `ParallelUpdater` does not manage the current device, so the backward process of `cupy.elementwise()` fails. 

The reproduction code below. Note F.mean_absolute_error() contains cupy.elementwise().

```example.py
import chainer
import chainer.links as L
import chainer.functions as F

from chainer.optimizers import Adam
from chainer.iterators import SerialIterator
from chainer.training import Trainer, ParallelUpdater

import numpy as np


class Model(chainer.Chain):
    def __init__(self):
        super().__init__()
        with self.init_scope():
            self.l1 = L.Linear(5, 5)

    def __call__(self, x):
        s = list(x.shape)
        s.pop()
        loss = F.mean_absolute_error(self.l1(x), x.reshape(s))
        return loss


class Dataset(chainer.dataset.DatasetMixin):
    def __init__(self):
        super().__init__()

    def __len__(self):
        return 1024

    def get_example(self, i):
        return np.arange(5, dtype=np.float32).reshape((5, 1))

model = Model()
optimizer = Adam()
optimizer.setup(model)

train_dataset = Dataset()
train_iterator = SerialIterator(train_dataset, 2)

updater = ParallelUpdater(train_iterator, optimizer, devices={"main": 0, "second": 1})
trainer = Trainer(updater)
trainer.run()
```
```traceback.py
Exception in main training loop: Array device must be same as the current device: array device = 1 while current = 0
Traceback (most recent call last):
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/training/trainer.py", line 306, in run
    update()
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/training/updaters/standard_updater.py", line 149, in update
    self.update_core()
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/training/updaters/parallel_updater.py", line 136, in update_core
    loss.backward(loss_scale=self.loss_scale)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/variable.py", line 963, in backward
    self._backward_main(retain_grad, loss_scale)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/variable.py", line 1092, in _backward_main
    target_input_indexes, out_grad, in_grad)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/function_node.py", line 548, in backward_accumulate
    gxs = self.backward(target_input_indexes, grad_outputs)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/loss/mean_absolute_error.py", line 37, in backward
    gx0 = coeff * cuda.get_array_module(gy.data).sign(self.diff)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/cupy/core/fusion.py", line 713, in __call__
    return self._cupy_op(*args, **kwargs)
  File "cupy/core/elementwise.pxi", line 786, in cupy.core.core.ufunc.__call__
  File "cupy/core/elementwise.pxi", line 135, in cupy.core.core._preprocess_args
Will finalize trainer extensions and updater before reraising the exception.
Traceback (most recent call last):
  File "b", line 44, in <module>
    trainer.run()
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/training/trainer.py", line 320, in run
    six.reraise(*sys.exc_info())
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/six.py", line 693, in reraise
    raise value
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/training/trainer.py", line 306, in run
    update()
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/training/updaters/standard_updater.py", line 149, in update
    self.update_core()
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/training/updaters/parallel_updater.py", line 136, in update_core
    loss.backward(loss_scale=self.loss_scale)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/variable.py", line 963, in backward
    self._backward_main(retain_grad, loss_scale)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/variable.py", line 1092, in _backward_main
    target_input_indexes, out_grad, in_grad)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/function_node.py", line 548, in backward_accumulate
    gxs = self.backward(target_input_indexes, grad_outputs)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/loss/mean_absolute_error.py", line 37, in backward
    gx0 = coeff * cuda.get_array_module(gy.data).sign(self.diff)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/cupy/core/fusion.py", line 713, in __call__
    return self._cupy_op(*args, **kwargs)
  File "cupy/core/elementwise.pxi", line 786, in cupy.core.core.ufunc.__call__
  File "cupy/core/elementwise.pxi", line 135, in cupy.core.core._preprocess_args
ValueError: Array device must be same as the current device: array device = 1 while current = 0
```


Discussion:
I ran into this as well attempting to adapt the RNN example in text_classification to use multiple gpus using `ParallelUpdater`: https://github.com/chainer/chainer/tree/master/examples/text_classification

Resolving patch:
diff --git a/chainer/variable.py b/chainer/variable.py
--- a/chainer/variable.py
+++ b/chainer/variable.py
@@ -1051,7 +1051,7 @@ def set_grad(node, value):
                 hooks.update(func.local_function_hooks)
             hooks = hooks.values()  # avoid six for performance
 
-            cuda.get_device_from_array(*in_data).use()
+            cuda.get_device_from_array(*(in_data + out_grad_data)).use()
             for hook in hooks:
                 hook.backward_preprocess(func, in_data, out_grad_data)
 


Test patch:
diff --git a/tests/chainer_tests/functions_tests/loss_tests/test_absolute_error.py b/tests/chainer_tests/functions_tests/loss_tests/test_absolute_error.py
--- a/tests/chainer_tests/functions_tests/loss_tests/test_absolute_error.py
+++ b/tests/chainer_tests/functions_tests/loss_tests/test_absolute_error.py
@@ -84,5 +84,16 @@ def test_double_backward_gpu(self):
             cuda.to_gpu(self.x0), cuda.to_gpu(self.x1), cuda.to_gpu(self.gy),
             cuda.to_gpu(self.ggx0), cuda.to_gpu(self.ggx1))
 
+    # test for #4669
+    @attr.multi_gpu(2)
+    def test_backward_non_default_gpu(self):
+        x0 = chainer.Variable(cuda.to_gpu(self.x0, 1))
+        x1 = chainer.Variable(cuda.to_gpu(self.x1, 1))
+        gy = cuda.to_gpu(self.gy, 1)
+        with cuda.get_device_from_id(0):
+            y = functions.absolute_error(x0, x1)
+            y.grad = gy
+            y.backward()
+
 
 testing.run_module(__name__, __file__)
diff --git a/tests/chainer_tests/functions_tests/loss_tests/test_mean_absolute_error.py b/tests/chainer_tests/functions_tests/loss_tests/test_mean_absolute_error.py
--- a/tests/chainer_tests/functions_tests/loss_tests/test_mean_absolute_error.py
+++ b/tests/chainer_tests/functions_tests/loss_tests/test_mean_absolute_error.py
@@ -78,5 +78,16 @@ def test_double_backward_gpu(self):
                                    cuda.to_gpu(self.ggx0),
                                    cuda.to_gpu(self.ggx1))
 
+    # test for #4669
+    @attr.multi_gpu(2)
+    def test_backward_non_default_gpu(self):
+        x0 = chainer.Variable(cuda.to_gpu(self.x0, 1))
+        x1 = chainer.Variable(cuda.to_gpu(self.x1, 1))
+        gy = cuda.to_gpu(self.gy, 1)
+        with cuda.get_device_from_id(0):
+            y = functions.mean_absolute_error(x0, x1)
+            y.grad = gy
+            y.backward()
+
 
 testing.run_module(__name__, __file__)