Repository: chainer/chainer

Issue:
Handle None as FunctionNode input with ChainerX
Required in seq2seq example.


Discussion:

Resolving patch:
diff --git a/chainer/function_node.py b/chainer/function_node.py
--- a/chainer/function_node.py
+++ b/chainer/function_node.py
@@ -418,20 +418,23 @@ def _chainerx_apply_fallback_preprocess(self, in_data, inputs):
         in_data = []
         device = None
         for data, x in six.moves.zip(chainerx_in_data, inputs):
-            # Use the cached fallback arrays as inputs if they exist.
-            x_is_variable = isinstance(x, variable.Variable)
-            if x_is_variable and x._chainerx_fallback_array is not None:
-                fallback_data = x._chainerx_fallback_array
-                if device is None:
-                    device = x.device
+            if data is None:
+                fallback_data = None
             else:
-                fallback_data = backend.from_chainerx(data)
-                if device is None:
-                    device = backend.ChainerxDevice(data.device)
+                # Use the cached fallback arrays as inputs if they exist.
+                x_is_variable = isinstance(x, variable.Variable)
+                if x_is_variable and x._chainerx_fallback_array is not None:
+                    fallback_data = x._chainerx_fallback_array
+                    if device is None:
+                        device = x.device
+                else:
+                    fallback_data = backend.from_chainerx(data)
+                    if device is None:
+                        device = backend.ChainerxDevice(data.device)
 
-                # Update the fallback cache if possible.
-                if x_is_variable:
-                    x._chainerx_fallback_array = fallback_data
+                    # Update the fallback cache if possible.
+                    if x_is_variable:
+                        x._chainerx_fallback_array = fallback_data
 
             in_data.append(fallback_data)
 
@@ -454,8 +457,9 @@ def _chainerx_apply_fallback_postprocess(
             [] if self._output_indexes_to_retain is None
             else self._output_indexes_to_retain)
 
-        self.inputs = tuple(
-            [variable._ChainerxVariableNodeProps(x) for x in inputs])
+        self.inputs = tuple([
+            None if x is None
+            else variable._ChainerxVariableNodeProps(x) for x in inputs])
 
         ret = tuple([
             _to_variable_with_chainerx_fallback_array(
@@ -717,7 +721,8 @@ def _backward_chainerx(self, target_input_indexes, grad_outputs,
             for a in grad_outputs])
 
         self._chainerx_retained_inputs = tuple([
-            variable.Variable(
+            None if array is None
+            else variable.Variable(
                 array, requires_grad=array.is_backprop_required())
             for array in retained_inputs])
         self._chainerx_retained_outputs = tuple([


Test patch:
diff --git a/tests/chainer_tests/test_backprop.py b/tests/chainer_tests/test_backprop.py
--- a/tests/chainer_tests/test_backprop.py
+++ b/tests/chainer_tests/test_backprop.py
@@ -240,10 +240,9 @@ def test_backprop(self, backend_config):
         {'use_cuda': True, 'cuda_device': 0},
         {'use_cuda': True, 'cuda_device': 1},
         # ChainerX
-        # TODO(niboshi): Fix ChainerX for None inputs
-        # {'use_chainerx': True, 'chainerx_device': 'native:0'},
-        # {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
-        # {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
+        {'use_chainerx': True, 'chainerx_device': 'native:0'},
+        {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
+        {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
     ])
 class TestFunctionInputNone(unittest.TestCase):