Repository: chainer/chainer

Issue:
chainer.functions.permutate() is incompatible to ParallelUpdater
Exactly same with #4669 except the problematic function is `chainer.functions.permutate()` here.
Maybe there are more and more functions which is incompatible to `ParallelUpdater`.

```test.py
import chainer
import chainer.functions as F

from chainer.optimizers import Adam
from chainer.iterators import SerialIterator
from chainer.training import Trainer, ParallelUpdater
from chainer.backends import cuda

import numpy as np


class Model(chainer.Chain):
    def __init__(self):
        super().__init__()

    def __call__(self, x):
        indices = cuda.get_array_module(x).asarray((1, 0))
        h = F.permutate(x, indices, 1, True)
        return F.mean(h)


class Dataset(chainer.dataset.DatasetMixin):
    def __init__(self):
        super().__init__()

    def __len__(self):
        return 1024

    def get_example(self, i):
        return np.arange(2, dtype=np.float32)


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, (1, "iteration"))
trainer.run()
```
```trace.py
/home/ubuntu/.python-venv/lib/python3.5/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
  from ._conv import register_converters as _register_converters
Exception in main training loop: Array device must be same as the current device: array device = 0 while current = 1
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 128, in update_core
    loss = loss_func(in_arrays)
  File "for_chainer_issue_5.py", line 18, in __call__
    h = F.permutate(x, indices, 1, True)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 134, in permutate
    y, = Permutate(axis, inv).apply((x, indices))
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/function_node.py", line 258, in apply
    outputs = self.forward(in_data)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 74, in forward
    return self._permutate(x, inds, self.inv),
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 63, in _permutate
    indices = _inverse_indices(indices)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 35, in _inverse_indices
    )(indices, r)
  File "cupy/core/elementwise.pxi", line 551, in cupy.core.core.ElementwiseKernel.__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 "for_chainer_issue_5.py", line 42, 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 128, in update_core
    loss = loss_func(in_arrays)
  File "for_chainer_issue_5.py", line 18, in __call__
    h = F.permutate(x, indices, 1, True)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 134, in permutate
    y, = Permutate(axis, inv).apply((x, indices))
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/function_node.py", line 258, in apply
    outputs = self.forward(in_data)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 74, in forward
    return self._permutate(x, inds, self.inv),
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 63, in _permutate
    indices = _inverse_indices(indices)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 35, in _inverse_indices
    )(indices, r)
  File "cupy/core/elementwise.pxi", line 551, in cupy.core.core.ElementwiseKernel.__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 = 0 while current = 1
```


Discussion:
Other version
```test.py
import chainer
import chainer.functions as F

from chainer.optimizers import Adam
from chainer.iterators import SerialIterator
from chainer.training import Trainer, ParallelUpdater
from chainer.backends import cuda

import numpy as np


class Model(chainer.Chain):
    def __init__(self):
        super().__init__()

    def __call__(self, x):
        indices = cuda.get_array_module(x).asarray((1, 0))
        h = F.permutate(x, indices, 1, False)
        return F.mean(h)


class Dataset(chainer.dataset.DatasetMixin):
    def __init__(self):
        super().__init__()

    def __len__(self):
        return 1024

    def get_example(self, i):
        return np.arange(2, dtype=np.float32)


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, (1, "iteration"))
trainer.run()
```
```trace.py
Exception in main training loop: Array device must be same as the current device: array device = 0 while current = 1
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 128, in update_core
    loss = loss_func(in_arrays)
  File "for_chainer_issue_5.py", line 18, in __call__
    h = F.permutate(x, indices, 1)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 134, in permutate
    y, = Permutate(axis, inv).apply((x, indices))
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/function_node.py", line 258, in apply
    outputs = self.forward(in_data)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 74, in forward
    return self._permutate(x, inds, self.inv),
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 65, in _permutate
    return x[((slice(None),) * self.axis) + (indices,)]
  File "cupy/core/core.pyx", line 1500, in cupy.core.core.ndarray.__getitem__
  File "cupy/core/core.pyx", line 683, in cupy.core.core.ndarray.take
  File "cupy/core/core.pyx", line 691, in cupy.core.core.ndarray.take
  File "cupy/core/core.pyx", line 3076, in cupy.core.core._take
  File "cupy/core/elementwise.pxi", line 551, in cupy.core.core.ElementwiseKernel.__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 "for_chainer_issue_5.py", line 42, 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 128, in update_core
    loss = loss_func(in_arrays)
  File "for_chainer_issue_5.py", line 18, in __call__
    h = F.permutate(x, indices, 1)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 134, in permutate
    y, = Permutate(axis, inv).apply((x, indices))
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/function_node.py", line 258, in apply
    outputs = self.forward(in_data)
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 74, in forward
    return self._permutate(x, inds, self.inv),
  File "/home/ubuntu/.python-venv/lib/python3.5/site-packages/chainer/functions/array/permutate.py", line 65, in _permutate
    return x[((slice(None),) * self.axis) + (indices,)]
  File "cupy/core/core.pyx", line 1500, in cupy.core.core.ndarray.__getitem__
  File "cupy/core/core.pyx", line 683, in cupy.core.core.ndarray.take
  File "cupy/core/core.pyx", line 691, in cupy.core.core.ndarray.take
  File "cupy/core/core.pyx", line 3076, in cupy.core.core._take
  File "cupy/core/elementwise.pxi", line 551, in cupy.core.core.ElementwiseKernel.__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 = 0 while current = 1
```
I think chainer had better add a test for chainer.functions w/ ParallelUpdater.

Resolving patch:
diff --git a/chainer/training/updaters/parallel_updater.py b/chainer/training/updaters/parallel_updater.py
--- a/chainer/training/updaters/parallel_updater.py
+++ b/chainer/training/updaters/parallel_updater.py
@@ -2,6 +2,7 @@
 
 import six
 
+from chainer.backends import cuda
 from chainer.dataset import convert
 from chainer import function
 from chainer.training.updaters import standard_updater
@@ -120,12 +121,16 @@ def update_core(self):
             loss_func = self.loss_func or model
 
             with function.force_backprop_mode():
-                if isinstance(in_arrays, tuple):
-                    loss = loss_func(*in_arrays)
-                elif isinstance(in_arrays, dict):
-                    loss = loss_func(**in_arrays)
-                else:
-                    loss = loss_func(in_arrays)
+                dev_id = self._devices[model_key]
+                dev_id = dev_id if 0 <= dev_id else None
+                with cuda.get_device_from_id(dev_id):
+                    if isinstance(in_arrays, tuple):
+                        loss = loss_func(*in_arrays)
+                    elif isinstance(in_arrays, dict):
+                        loss = loss_func(**in_arrays)
+                    else:
+                        loss = loss_func(in_arrays)
+
             losses.append(loss)
 
         # For _uninitialized_params


Test patch: