Merge 88a6cfd into 728d941

chainer/functions/math/average.py

@@ -1,8 +1,83 @@
 import six
 
+import chainer
+from chainer import function_node
 from chainer.functions.array import broadcast
 from chainer.functions.array import reshape
 from chainer.functions.math import sum as sum_mod
+from chainer import utils
+from chainer.utils import type_check
+
+
+class Mean(function_node.FunctionNode):
+    """Mean of array elements over a given axis."""
+
+    def __init__(self, axis, keepdims):
+        if axis is None:
+            self.axis = None
+        elif isinstance(axis, six.integer_types):
+            self.axis = (axis,)
+        elif isinstance(axis, tuple) and all(
+                isinstance(a, six.integer_types) 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')
+
+        self.keepdims = keepdims
+
+    def check_type_forward(self, in_types):
+        type_check._argname(in_types, ('x',))
+        type_check.expect(in_types[0].dtype.kind == 'f')
+
+        if self.axis is not None:
+            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,
+                    )
+
+    # TODO(kataoka): override `forward_chainerx` if `chainerx.mean` does not
+    # overflow for large float16 inputs
+
+    def forward(self, inputs):
+        x, = inputs
+        if self.axis is None:
+            self.multiplier = 1.0 / x.size
+        else:
+            divider = 1
+            for axis in self.axis:
+                divider *= x.shape[axis]
+            self.multiplier = 1.0 / divider
+        ret = utils.force_array(
+            x.mean(axis=self.axis, keepdims=self.keepdims))
+        return ret,
+
+    def backward(self, indexes, grad_outputs):
+        gy, = grad_outputs
+        gy *= self.multiplier
+        ndim = len(self.inputs[0].shape)
+        if not (ndim == 0 or self.axis is None or self.keepdims):
+            actual_axis = [
+                axis if axis >= 0 else axis + ndim
+                for axis in self.axis]
+            shape = list(gy.shape)
+            for axis in sorted(actual_axis):
+                shape.insert(axis, 1)
+            gy = chainer.functions.reshape(gy, shape)
+        return chainer.functions.broadcast_to(gy, self.inputs[0].shape),
+
+
+# TODO(kataoka): consider making the function public.
+def _mean(x, axis, keepdims=False):
+    y, = Mean(axis, keepdims).apply((x,))
+    return y
 
 
 def average(x, axis=None, weights=None, keepdims=False):
@@ -26,6 +101,8 @@ def average(x, axis=None, weights=None, keepdims=False):
         ~chainer.Variable: Output variable.
 
     """
+    if weights is None:
+        return _mean(x, axis, keepdims)
     if axis is None:
         pass
     elif isinstance(axis, tuple):
@@ -40,28 +117,17 @@ def average(x, axis=None, weights=None, keepdims=False):
             axis += x.ndim
         axis = (axis,)
 
-    if weights is not None:
-        if axis is not None and len(axis) > 1:
-            raise ValueError(
-                'tuple axis is not supported when weights is given')
-        divider = sum_mod.sum(weights)
-        if axis is not None:
-            w_shape = [d if i in axis else 1 for i, d in enumerate(x.shape)]
-            weights = broadcast.broadcast_to(
-                reshape.reshape(weights, w_shape), x.shape)
-
-        x = x * weights
-    else:
-        if axis is None:
-            divider = x.size
-        else:
-            divider = 1
-            for a in axis:
-                divider *= x.shape[a]
+    if axis is not None and len(axis) > 1:
+        raise ValueError(
+            'tuple axis is not supported when weights is given')
+    divider = sum_mod.sum(weights)
+    if axis is not None:
+        w_shape = [d if i in axis else 1 for i, d in enumerate(x.shape)]
+        weights = broadcast.broadcast_to(
+            reshape.reshape(weights, w_shape), x.shape)
+
+    x = x * weights
 
     x_sum = sum_mod.sum(x, axis, keepdims)
-    if weights is not None:
-        # We do not need to call broadcast when weights is None because
-        # divider here is not a Variable but a scalar
-        divider = broadcast.broadcast_to(divider, x_sum.shape)
+    divider = broadcast.broadcast_to(divider, x_sum.shape)
     return x_sum / divider

tests/chainer_tests/functions_tests/math_tests/test_average.py

@@ -5,6 +5,7 @@
 
 from chainer import functions
 from chainer import testing
+from chainer.testing import attr
 from chainer import utils
 
 
@@ -114,6 +115,58 @@ def forward_expected(self, inputs):
         return y_expect,
 
 
+@testing.parameterize(*(
+    testing.product({
+        'shape': [(30, 20, 40)],
+        'axis': [None, 0, 1, 2, -1, (0, 1), (1, -1)],
+        'dtype': [numpy.float16],
+        'use_weights': [False],  # np.average overflows when `weights` is used
+        'keepdims': [True, False],
+    })
+))
+@testing.inject_backend_tests(
+    None,
+    # CPU tests
+    [
+        {},
+    ]
+    # GPU tests
+    + testing.product({
+        'use_cuda': [True],
+        'cuda_device': [0, 1],
+    })
+    # ChainerX tests
+    + testing.product({
+        'use_chainerx': [True],
+        'chainerx_device': ['native:0', 'cuda:0', 'cuda:1'],
+    }))
+@attr.slow
+@testing.with_requires('numpy>=1.12')  # NumPy #8222
+class TestAverageOverflowingSum(testing.FunctionTestCase):
+
+    def setUp(self):
+        self.check_forward_options.update({'atol': 1e-2, 'rtol': 2e-3})
+        self.check_backward_options.update({'atol': 1e-2, 'rtol': 1e-2})
+        self.check_double_backward_options.update({'atol': 1e-2, 'rtol': 1e-2})
+
+    def generate_inputs(self):
+        x = numpy.random.uniform(3000, 7000, self.shape).astype(self.dtype)
+        return x,
+
+    def forward(self, inputs, device):
+        x, = inputs
+        y = functions.average(
+            x, self.axis, keepdims=self.keepdims)
+        return y,
+
+    def forward_expected(self, inputs):
+        x, = inputs
+        y_expect = numpy.mean(
+            x.astype(numpy.float64), self.axis, keepdims=self.keepdims
+        ).astype(self.dtype)
+        return utils.force_array(y_expect),
+
+
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float16, numpy.float32, numpy.float64],
 }))