Repository: chainer/chainer

Issue:
split_axis.backward fails on incomplete gradients
When there is a None in the grad_outputs, split_axis fails to backprop the incomplete gradients.



Discussion:
minimum reproducing example:

``` python
import chainer
from chainer import cuda
import chainer.functions as F
from chainer import optimizers
import numpy as np

xs = chainer.Variable(np.array([i for i in range(100)],dtype=np.float32))
_, smaller, _ = F.split_axis(xs, [24,47], axis=0)
smaller.backward()
```

result:

```
$ python split.py 
Traceback (most recent call last):
  File "split.py", line 11, in <module>
    smaller.backward()
  File "build/bdist.linux-x86_64/egg/chainer/variable.py", line 189, in backward
  File "build/bdist.linux-x86_64/egg/chainer/function.py", line 317, in backward
  File "build/bdist.linux-x86_64/egg/chainer/functions/split_axis.py", line 88, in backward_cpu
ValueError: zero-dimensional arrays cannot be concatenated
```


Resolving patch:
diff --git a/chainer/functions/split_axis.py b/chainer/functions/split_axis.py
--- a/chainer/functions/split_axis.py
+++ b/chainer/functions/split_axis.py
@@ -1,6 +1,7 @@
 import collections
 
 import numpy
+import six
 
 from chainer import cuda
 from chainer import function
@@ -85,14 +86,25 @@ def forward_gpu(self, x):
         return tuple(ys)
 
     def backward_cpu(self, x, gys):
-        return numpy.concatenate(gys, axis=self.axis),
+        if any(gy is None for gy in gys):
+            gx = numpy.zeros_like(x[0])
+            gxs = numpy.split(gx, self.indices_or_sections, self.axis)
+            for gxi, gy in six.moves.zip(gxs, gys):
+                if gy is None:
+                    continue
+                gxi[:] = gy
+            return gx,
+        else:
+            return numpy.concatenate(gys, axis=self.axis),
 
     def backward_gpu(self, x, gys):
-        gx = cuda.empty_like(x[0])
+        gx = cuda.zeros_like(x[0])
         coffset = 0
         kernel = cuda.elementwise(
             _args, 'COPY(x[idx] = y[i])', 'split_bwd', preamble=_preamble)
         for gy in gys:
+            if gy is None:
+                continue
             cdimy = gy.shape[self.axis]
             if cdimy != 0:
                 kernel(gy, gx, cdimy, self.cdimx, self.rdim, coffset)


Test patch:
diff --git a/tests/functions_tests/test_split_axis.py b/tests/functions_tests/test_split_axis.py
--- a/tests/functions_tests/test_split_axis.py
+++ b/tests/functions_tests/test_split_axis.py
@@ -93,4 +93,30 @@ def setUp(self):
         self.axis = 0
 
 
+class TestSplitAxisNone(unittest.TestCase):
+
+    def setUp(self):
+        self.x = numpy.array([1, 2], dtype=numpy.float32)
+        self.ys_section = [1]
+        self.axis = 0
+
+    def check_backward(self, x_data, indices_or_sections, axis):
+        x = chainer.Variable(x_data)
+        ys = functions.split_axis(x, indices_or_sections, axis)
+        # Only set ys[0]
+        ys[0].grad = ys[0].data
+        ys[0].backward()
+
+        gx = numpy.array([1, 0])
+        gradient_check.assert_allclose(gx, x.grad, atol=0, rtol=0)
+
+    def test_backward_cpu(self):
+        self.check_backward(self.x, self.ys_section, axis=self.axis)
+
+    @attr.gpu
+    def test_backward_gpu(self):
+        self.check_backward(
+            cuda.to_gpu(self.x), self.ys_section, axis=self.axis)
+
+
 testing.run_module(__name__, __file__)