Repository: chainer/chainer

Issue:
Cannot deepcopy a link after to_device('native')
I used chainer and chainerx of the current master https://github.com/chainer/chainer/commit/e413741d8134b578b43bf1f05c1de3b18bd2d9a6

Both `mode='copy'` and `mode='init'` fail.

Code
```
import chainer.links as L


model = L.Linear(2, 3)
model.to_device('native')

model2 = model.copy(mode='copy')
```

Output
```
$ python check_chainerx_deepcopy.py
Traceback (most recent call last):
  File "check_chainerx_deepcopy.py", line 7, in <module>
    model2 = model.copy(mode='copy')
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/site-packages/chainer/link.py", line 388, in copy
    return copy.deepcopy(self)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/copy.py", line 182, in deepcopy
    y = _reconstruct(x, rv, 1, memo)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/copy.py", line 297, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/copy.py", line 155, in deepcopy
    y = copier(x, memo)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/copy.py", line 243, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/copy.py", line 182, in deepcopy
    y = _reconstruct(x, rv, 1, memo)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/copy.py", line 292, in _reconstruct
    y = callable(*args)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/site-packages/chainer/variable.py", line 1813, in _recover_parameter
    p.array = data
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/site-packages/chainer/variable.py", line 819, in array
    self._node._update_data_info(d)
  File "/home/fujita/.local/share/virtualenvs/chainerrl_20180827-9TPAa4mI/lib/python3.5/site-packages/chainer/variable.py", line 415, in _update_data_info
    self.dtype = d.dtype
RuntimeError: Unable to cast from non-held to held instance (T& to Holder<T>) (compile in debug mode for type information)
```


Discussion:
I assume the cause is that `chainerx.Device` is not picklable (according to the error in my environment):
```
TypeError: can't pickle chainerx.Device objects
```

The difficult point is that a `chainerx.Device` is associated with a `chainerx.Context`, which is never (or extremely difficult to make) picklable.

One idea: First give each context a unique name (specified by user. Empty for the default context.). When pickling, the device ID and the context name are saved. When unpickling, a context with the saved name is looked for and if found, restore `chainerx.Device` from that context. If not, an error will be raised.

Resolving patch:
diff --git a/chainer/variable.py b/chainer/variable.py
--- a/chainer/variable.py
+++ b/chainer/variable.py
@@ -428,9 +428,11 @@ def _check_old_style_gradient(self):
                 self._old_style_grad_generator)
 
 
-def _create_variable(data, name, grad, requires_grad):
-    return Variable(
+def _create_variable(data, name, grad, requires_grad, device):
+    var = Variable(
         data, name=name, grad=grad, requires_grad=requires_grad)
+    var.to_device(device)
+    return var
 
 
 class Variable(object):
@@ -555,8 +557,9 @@ def _copy_to(self, target):
         return target
 
     def __reduce__(self):
-        return _create_variable, (self.array, self.name, self.grad,
-                                  self._requires_grad)
+        args = (
+            self.array, self.name, self.grad, self._requires_grad, self.device)
+        return _create_variable, args
 
     def __repr__(self):
         return variable_repr(self)
@@ -1660,8 +1663,10 @@ def __copy__(self):
         return self._copy_to(Parameter())
 
     def __reduce__(self):
-        return _recover_parameter, (self.array, self.name, self.grad,
-                                    self.initializer, self.update_rule)
+        args = (
+            self.array, self.name, self.grad, self.initializer,
+            self.update_rule, self.device)
+        return _recover_parameter, args
 
     def to_cpu(self):
         return self.to_device(backend.CpuDevice())
@@ -1823,11 +1828,12 @@ def as_array(obj):
     return obj
 
 
-def _recover_parameter(data, name, grad, initializer, update_rule):
+def _recover_parameter(data, name, grad, initializer, update_rule, device):
     p = Parameter(initializer=initializer, name=name)
     p.array = data
     p.grad = grad
     p.update_rule = update_rule
+    p.to_device(device)
     return p
 
 
diff --git a/chainerx/__init__.py b/chainerx/__init__.py
--- a/chainerx/__init__.py
+++ b/chainerx/__init__.py
@@ -20,6 +20,8 @@
 
     from builtins import bool, int, float  # NOQA
 
+    from chainerx import _device  # NOQA
+
     from chainerx.creation.from_data import asanyarray  # NOQA
     from chainerx.creation.from_data import fromfile  # NOQA
     from chainerx.creation.from_data import fromfunction  # NOQA
diff --git a/chainerx/_device.py b/chainerx/_device.py
new file mode 100644
--- /dev/null
+++ b/chainerx/_device.py
@@ -0,0 +1,11 @@
+import chainerx
+
+
+def _recover_device(backend_name, device_index):
+    # Recovers the device instance.
+    # This function is used together with chainerx.Device.__reduce__.
+    # TODO(niboshi): Save the context name and lookup the context with it.
+    context = chainerx.get_default_context()
+    backend = context.get_backend(backend_name)
+    device = backend.get_device(device_index)
+    return device


Test patch:
diff --git a/tests/chainerx_tests/unit_tests/test_array.py b/tests/chainerx_tests/unit_tests/test_array.py
--- a/tests/chainerx_tests/unit_tests/test_array.py
+++ b/tests/chainerx_tests/unit_tests/test_array.py
@@ -1,4 +1,6 @@
+import copy
 import math
+import pickle
 
 import numpy
 import pytest
@@ -697,3 +699,33 @@ def test_asarray_to_numpy_identity(device, slice1, slice2):
     z = chainerx.to_numpy(y)
     chainerx.testing.assert_array_equal_ex(x, y)
     chainerx.testing.assert_array_equal_ex(x, z, strides_check=False)
+
+
+# TODO(niboshi): Add pickle test involving context destruction and re-creation
+@pytest.mark.parametrize_device(['native:0', 'native:1', 'cuda:0'])
+def test_array_pickle(device):
+    arr = chainerx.array([1, 2], chainerx.float32, device=device)
+    s = pickle.dumps(arr)
+    del arr
+
+    arr2 = pickle.loads(s)
+    assert isinstance(arr2, chainerx.ndarray)
+    assert arr2.device is device
+    assert arr2.dtype == chainerx.float32
+    chainerx.testing.assert_array_equal(
+        arr2,
+        chainerx.array([1, 2], chainerx.float32))
+
+
+# TODO(niboshi): Add deepcopy test with arbitrary context
+@pytest.mark.parametrize_device(['native:0', 'native:1', 'cuda:0'])
+def test_array_deepcopy(device):
+    arr = chainerx.array([1, 2], chainerx.float32, device=device)
+    arr2 = copy.deepcopy(arr)
+
+    assert isinstance(arr2, chainerx.ndarray)
+    assert arr2.device is device
+    assert arr2.dtype == chainerx.float32
+    chainerx.testing.assert_array_equal(
+        arr2,
+        chainerx.array([1, 2], chainerx.float32))
diff --git a/tests/chainerx_tests/unit_tests/test_device.py b/tests/chainerx_tests/unit_tests/test_device.py
--- a/tests/chainerx_tests/unit_tests/test_device.py
+++ b/tests/chainerx_tests/unit_tests/test_device.py
@@ -1,3 +1,6 @@
+import copy
+import pickle
+
 import pytest
 
 import chainerx
@@ -133,3 +136,18 @@ def test_using_device_with_name(device_instance1, device_instance2):
     with chainerx.using_device(device2.backend.name, device2.index) as scope:
         assert chainerx.get_default_device() == device2
         assert scope.device is device2
+
+
+# TODO(niboshi): Add pickle test involving context destruction and re-creation
+@pytest.mark.parametrize_device(['native:0', 'native:1', 'cuda:0'])
+def test_device_pickle(device):
+    s = pickle.dumps(device)
+    device2 = pickle.loads(s)
+    assert device is device2
+
+
+# TODO(niboshi): Add deepcopy test with arbitrary context
+@pytest.mark.parametrize_device(['native:0', 'native:1', 'cuda:0'])
+def test_device_deepcopy(device):
+    device2 = copy.deepcopy(device)
+    assert device is device2