Repository: chainer/chainer

Issue:
GPU computation in L.NStepLSTM may fail if not on the current device
`xp.array` outside `FunctionNode` leads to this kind of bugs.

* Conditions
```
Platform: Linux-4.4.0-103-generic-x86_64-with-debian-stretch-sid
Chainer: 5.0.0b4
NumPy: 1.16.0.dev0+91454d6
CuPy:
  CuPy Version          : 5.0.0b4
  CUDA Root             : /usr/local/cuda
  CUDA Build Version    : 9000
  CUDA Driver Version   : 9000
  CUDA Runtime Version  : 9000
  cuDNN Build Version   : 7104
  cuDNN Version         : 7104
  NCCL Build Version    : 2005
```
* Code to reproduce
```
>>> lstm = chainer.links.NStepLSTM(1, 9, 9, 0)
>>> lstm.to_gpu(1)
<chainer.links.connection.n_step_lstm.NStepLSTM object at 0x2baa76604c88>
>>> x = chainer.Variable(numpy.random.rand(3, 9).astype('f'))
>>> x.to_gpu(1)
```
* Error messages, stack traces, or logs
```
>>> lstm(None, None, [x])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/kataoka/chainer/chainer/link.py", line 204, in __call__
    return forward(*args, **kwargs)
  File "/home/kataoka/chainer/chainer/links/connection/n_step_lstm.py", line 70, in forward
    (hy, cy), ys = self._call([hx, cx], xs, **kwargs)
  File "/home/kataoka/chainer/chainer/links/connection/n_step_rnn.py", line 209, in _call
    for h in result[:-1]]
  File "/home/kataoka/chainer/chainer/links/connection/n_step_rnn.py", line 209, in <listcomp>
    for h in result[:-1]]
  File "/home/kataoka/chainer/chainer/functions/array/permutate.py", line 134, in permutate
    y, = Permutate(axis, inv).apply((x, indices))
  File "/home/kataoka/chainer/chainer/function_node.py", line 257, in apply
    outputs = self.forward(in_data)
  File "/home/kataoka/chainer/chainer/functions/array/permutate.py", line 74, in forward
    return self._permutate(x, inds, self.inv),
  File "/home/kataoka/chainer/chainer/functions/array/permutate.py", line 63, in _permutate
    indices = _inverse_indices(indices)
  File "/home/kataoka/chainer/chainer/functions/array/permutate.py", line 35, in _inverse_indices
    )(indices, r)
  File "cupy/core/_kernel.pyx", line 513, in cupy.core._kernel.ElementwiseKernel.__call__
  File "cupy/core/_kernel.pyx", line 74, in cupy.core._kernel._preprocess_args
ValueError: Array device must be same as the current device: array device = 0 while current = 1
```


Discussion:

Resolving patch:
diff --git a/chainer/functions/array/permutate.py b/chainer/functions/array/permutate.py
--- a/chainer/functions/array/permutate.py
+++ b/chainer/functions/array/permutate.py
@@ -40,24 +40,19 @@ class Permutate(function_node.FunctionNode):
 
     """Permutate function."""
 
-    def __init__(self, axis=0, inv=False):
+    def __init__(self, indices, axis, inv):
+        self.indices = indices
         self.axis = axis
         self.inv = inv
 
     def check_type_forward(self, in_types):
-        type_check.argname(in_types, ('x', 'indices'))
-        x_type, ind_type = in_types
+        type_check.argname(in_types, ('x',))
+        x_type, = in_types
         if self.axis < 0:
             type_check.expect(x_type.ndim >= -self.axis)
         else:
             type_check.expect(x_type.ndim > self.axis)
 
-        type_check.expect(
-            ind_type.dtype.kind == 'i',
-            ind_type.ndim == 1,
-            x_type.shape[self.axis] == ind_type.shape[0],
-        )
-
     def _permutate(self, x, indices, inv):
         if inv:
             indices = _inverse_indices(indices)
@@ -65,8 +60,8 @@ def _permutate(self, x, indices, inv):
         return x[((slice(None),) * self.axis) + (indices,)]
 
     def forward(self, inputs):
-        self.retain_inputs((1,))
-        x, inds = inputs
+        x, = inputs
+        inds = self.indices
 
         if chainer.is_debug():
             _check_indices(inds)
@@ -74,10 +69,10 @@ def forward(self, inputs):
         return self._permutate(x, inds, self.inv),
 
     def backward(self, indexes, grad_outputs):
-        inds = self.inputs[1]
         g, = grad_outputs
-        gx, = Permutate(self.axis, not self.inv).apply((g, inds.data))
-        return gx, None
+        inds = self.indices
+        gx, = Permutate(inds, self.axis, not self.inv).apply((g,))
+        return gx,
 
 
 def permutate(x, indices, axis=0, inv=False):
@@ -131,5 +126,10 @@ def permutate(x, indices, axis=0, inv=False):
                [5., 4.]], dtype=float32)
 
     """
-    y, = Permutate(axis, inv).apply((x, indices))
+    if indices.dtype.kind != 'i' or indices.ndim != 1:
+        raise ValueError(
+            'indices should be a one-dimensional int array')
+    if isinstance(indices, chainer.Variable):
+        indices = indices.array
+    y, = Permutate(indices, axis, inv).apply((x,))
     return y
diff --git a/chainer/links/connection/n_step_rnn.py b/chainer/links/connection/n_step_rnn.py
--- a/chainer/links/connection/n_step_rnn.py
+++ b/chainer/links/connection/n_step_rnn.py
@@ -186,9 +186,7 @@ def _call(self, hs, xs, **kwargs):
             argument.assert_kwargs_empty(kwargs)
 
         assert isinstance(xs, (list, tuple))
-        xp = cuda.get_array_module(*(list(hs) + list(xs)))
         indices = argsort_list_descent(xs)
-        indices_array = xp.array(indices)
 
         xs = permutate_list(xs, indices, inv=False)
         hxs = []
@@ -196,7 +194,7 @@ def _call(self, hs, xs, **kwargs):
             if hx is None:
                 hx = self.init_hx(xs)
             else:
-                hx = permutate.permutate(hx, indices_array, axis=1, inv=False)
+                hx = permutate.permutate(hx, indices, axis=1, inv=False)
             hxs.append(hx)
 
         trans_x = transpose_sequence.transpose_sequence(xs)
@@ -205,7 +203,7 @@ def _call(self, hs, xs, **kwargs):
                [self.ws, self.bs, trans_x]
         result = self.rnn(*args)
 
-        hys = [permutate.permutate(h, indices_array, axis=1, inv=True)
+        hys = [permutate.permutate(h, indices, axis=1, inv=True)
                for h in result[:-1]]
         trans_y = result[-1]
         ys = transpose_sequence.transpose_sequence(trans_y)


Test patch:
diff --git a/tests/chainer_tests/functions_tests/array_tests/test_permutate.py b/tests/chainer_tests/functions_tests/array_tests/test_permutate.py
--- a/tests/chainer_tests/functions_tests/array_tests/test_permutate.py
+++ b/tests/chainer_tests/functions_tests/array_tests/test_permutate.py
@@ -49,6 +49,10 @@ def test_forward_cpu(self):
     def test_forward_gpu(self):
         self.check_forward(cuda.to_gpu(self.x), cuda.to_gpu(self.indices))
 
+    @attr.gpu
+    def test_forward_mixed(self):
+        self.check_forward(cuda.to_gpu(self.x), self.indices)
+
     def check_backward(self, x_data, ind_data, g_data):
         def fun(x, ind):
             return functions.permutate(x, ind, self.axis, self.inv)
diff --git a/tests/chainer_tests/links_tests/connection_tests/test_n_step_lstm.py b/tests/chainer_tests/links_tests/connection_tests/test_n_step_lstm.py
--- a/tests/chainer_tests/links_tests/connection_tests/test_n_step_lstm.py
+++ b/tests/chainer_tests/links_tests/connection_tests/test_n_step_lstm.py
@@ -125,6 +125,18 @@ def test_forward_gpu_test(self):
                 cuda.to_gpu(self.c),
                 [cuda.to_gpu(x) for x in self.xs])
 
+    @attr.multi_gpu(2)
+    def test_forward_nonzero_gpu_test(self):
+        # Issue #5347
+        # to_gpu should work without setting the current device
+        self.rnn.to_gpu(1)
+        with chainer.using_config('use_cudnn', 'always'), \
+                chainer.using_config('train', False):
+            self.check_forward(
+                cuda.to_gpu(self.h, 1),
+                cuda.to_gpu(self.c, 1),
+                [cuda.to_gpu(x, 1) for x in self.xs])
+
     def check_backward(
             self, h_data, c_data, xs_data, gh_data, gc_data, gys_data):