Repository: chainer/chainer

Issue:
Add a function for summation along axes
`sum` function reduces all elements, while there are many cases that I want to take a summation over some axes. This is analogous to `numpy.sum` with the `axis` argument.

The implementation should contain a general function for GPU to compute a reduction over specified axes.



Discussion:
Summation along single axis is implemented in #337. Multiple axes is not supported yet


Resolving patch:
diff --git a/chainer/functions/math/sum.py b/chainer/functions/math/sum.py
--- a/chainer/functions/math/sum.py
+++ b/chainer/functions/math/sum.py
@@ -9,7 +9,17 @@ class Sum(function.Function):
     """Sum of array elements over a given axis."""
 
     def __init__(self, axis=None):
-        self.axis = axis
+        if axis is None:
+            self.axis = None
+        elif isinstance(axis, int):
+            self.axis = (axis,)
+        elif isinstance(axis, tuple) and all(isinstance(a, int) for a in axis):
+            if len(set(axis)) != len(axis):
+                raise ValueError('duplicate value in axis: ({})'.format(
+                    ', '.join(map(str, axis))))
+            self.axis = axis
+        else:
+            raise TypeError('None, int or tuple of int are required')
 
     def check_type_forward(self, in_types):
         type_check.expect(
@@ -18,9 +28,15 @@ def check_type_forward(self, in_types):
         )
 
         if self.axis is not None:
-            type_check.expect(
-                self.axis < in_types[0].ndim,
-            )
+            for axis in self.axis:
+                if axis >= 0:
+                    type_check.expect(
+                        axis < in_types[0].ndim,
+                    )
+                else:
+                    type_check.expect(
+                        -axis - 1 < in_types[0].ndim,
+                    )
 
     def forward(self, x):
         xp = cuda.get_array_module(*x)
@@ -33,7 +49,15 @@ def backward(self, x, gy):
         if self.axis is None:
             gx[:] = gy[0]
         else:
-            gx[:] = xp.expand_dims(gy[0], axis=self.axis)
+            gy = gy[0]
+            actual_axis = []
+            for axis in self.axis:
+                if axis < 0:
+                    axis = len(gx.shape) + axis
+                actual_axis.append(axis)
+            for axis in sorted(actual_axis):
+                gy = xp.expand_dims(gy, axis=axis)
+            gx[:] = gy
 
         return gx,
 
@@ -43,7 +67,7 @@ def sum(x, axis=None):
 
     Args:
         x (~chainer.Variable): Elements to sum.
-        axis (None or int): Axis which a sum is performed.
+        axis (None, int, or tuple of int): Axis which a sum is performed.
             The default (axis = None) is perform a sum over all the dimensions
             of the input array.
 


Test patch:
diff --git a/tests/chainer_tests/functions_tests/math_tests/test_sum.py b/tests/chainer_tests/functions_tests/math_tests/test_sum.py
--- a/tests/chainer_tests/functions_tests/math_tests/test_sum.py
+++ b/tests/chainer_tests/functions_tests/math_tests/test_sum.py
@@ -28,17 +28,67 @@ def check_forward(self, x_data, axis=None):
     def test_forward_cpu(self):
         self.check_forward(self.x)
 
+    @condition.retry(3)
+    def test_forward_axis_cpu(self):
         for i in range(self.x.ndim):
             self.check_forward(self.x, axis=i)
 
+    @condition.retry(3)
+    def test_forward_negative_axis_cpu(self):
+        self.check_forward(self.x, axis=-1)
+
+    @condition.retry(3)
+    def test_forward_multi_axis_cpu(self):
+        self.check_forward(self.x, axis=(0, 1))
+
+    @condition.retry(3)
+    def test_forward_multi_axis_invert_cpu(self):
+        self.check_forward(self.x, axis=(1, 0))
+
+    @condition.retry(3)
+    def test_forward_negative_multi_axis_cpu(self):
+        self.check_forward(self.x, axis=(0, -1))
+
+    @condition.retry(3)
+    def test_forward_negative_multi_axis_invert_cpu(self):
+        self.check_forward(self.x, axis=(-2, 0))
+
     @attr.gpu
     @condition.retry(3)
     def test_forward_gpu(self):
         self.check_forward(cuda.to_gpu(self.x))
 
+    @attr.gpu
+    @condition.retry(3)
+    def test_forward_axis_gpu(self):
         for i in range(self.x.ndim):
             self.check_forward(cuda.to_gpu(self.x), axis=i)
 
+    @attr.gpu
+    @condition.retry(3)
+    def test_forward_negative_axis_gpu(self):
+        self.check_forward(cuda.to_gpu(self.x), axis=-1)
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_forward_multi_axis_gpu(self):
+        self.check_forward(cuda.to_gpu(self.x), axis=(0, 1))
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_forward_multi_axis_invert_gpu(self):
+        self.check_forward(cuda.to_gpu(self.x), axis=(1, 0))
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_forward_negative_multi_axis_gpu(self):
+        self.check_forward(cuda.to_gpu(self.x), axis=(0, -1))
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_forward_negative_multi_axis_invert_gpu(self):
+        self.check_forward(cuda.to_gpu(self.x), axis=(-2, 0))
+
     def check_backward(self, x_data, y_grad, axis=None):
         x = chainer.Variable(x_data)
         y = functions.sum(x, axis=axis)
@@ -53,18 +103,94 @@ def check_backward(self, x_data, y_grad, axis=None):
     def test_backward_cpu(self):
         self.check_backward(self.x, self.gy)
 
+    @condition.retry(3)
+    def test_backward_axis_cpu(self):
         for i in range(self.x.ndim):
             gy = numpy.ones_like(self.x.sum(axis=i)) * self.gy
             self.check_backward(self.x, gy, axis=i)
 
+    @condition.retry(3)
+    def test_backward_negative_axis_cpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=-1)) * self.gy
+        self.check_backward(self.x, gy, axis=-1)
+
+    @condition.retry(3)
+    def test_backward_multi_axis_cpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(0, 1))) * self.gy
+        self.check_backward(self.x, gy, axis=(0, 1))
+
+    @condition.retry(3)
+    def test_backward_multi_axis_invert_cpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(1, 0))) * self.gy
+        self.check_backward(self.x, gy, axis=(1, 0))
+
+    @condition.retry(3)
+    def test_backward_negative_multi_axis_cpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(0, -1))) * self.gy
+        self.check_backward(self.x, gy, axis=(0, -1))
+
+    @condition.retry(3)
+    def test_backward_negative_multi_axis_invert_cpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(-2, 0))) * self.gy
+        self.check_backward(self.x, gy, axis=(-2, 0))
+
     @attr.gpu
     @condition.retry(3)
     def test_backward_gpu(self):
         self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.gy))
 
+    @attr.gpu
+    @condition.retry(3)
+    def test_backward_axis_gpu(self):
         for i in range(self.x.ndim):
             gy = numpy.ones_like(self.x.sum(axis=i)) * self.gy
             self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(gy), axis=i)
 
+    @attr.gpu
+    @condition.retry(3)
+    def test_backward_negative_axis_gpu(self):
+        for i in range(self.x.ndim):
+            gy = numpy.ones_like(self.x.sum(axis=-1)) * self.gy
+            self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(gy), axis=-1)
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_backward_multi_axis_gpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(0, 1))) * self.gy
+        self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(gy), axis=(0, 1))
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_backward_multi_axis_invert_gpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(1, 0))) * self.gy
+        self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(gy), axis=(1, 0))
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_backward_negative_multi_axis_gpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(0, -1))) * self.gy
+        self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(gy), axis=(0, -1))
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_backward_negative_multi_axis_invert_gpu(self):
+        gy = numpy.ones_like(self.x.sum(axis=(-2, 0))) * self.gy
+        self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(gy), axis=(-2, 0))
+
+    def test_invalid_axis_type(self):
+        with self.assertRaises(TypeError):
+            functions.Sum([0])
+
+    def test_invalid_axis_type_in_tuple(self):
+        with self.assertRaises(TypeError):
+            functions.Sum((1, 'x'))
+
+    def test_duplicate_axis(self):
+        with self.assertRaises(ValueError):
+            functions.Sum((0, 0))
+
+    def test_pos_neg_duplicate_axis(self):
+        with self.assertRaises(ValueError):
+            self.x.sum(axis=(1, -2))
 
 testing.run_module(__name__, __file__)