Repository: chainer/chainer

Issue:
Automatic broadcast for F.maximum, F.minimum
These functions should support automatic broadcast.


Discussion:

Resolving patch:
diff --git a/chainer/functions/array/where.py b/chainer/functions/array/where.py
--- a/chainer/functions/array/where.py
+++ b/chainer/functions/array/where.py
@@ -1,5 +1,6 @@
 import numpy
 
+import chainer
 from chainer.backends import cuda
 from chainer import function_node
 from chainer.utils import type_check
@@ -16,11 +17,12 @@ def check_type_forward(self, in_types):
         type_check.expect(
             c_type.dtype == numpy.bool_,
             x_type.dtype == y_type.dtype,
-            x_type.shape == c_type.shape,
-            y_type.shape == c_type.shape,
         )
+        type_check.expect_broadcast_shapes(
+            c_type.shape, x_type.shape, y_type.shape)
 
     def forward(self, inputs):
+        # may broadcast
         self.retain_inputs((0,))
         xp = cuda.get_array_module(*inputs)
         condition, x, y = inputs
@@ -30,16 +32,16 @@ def backward(self, indexes, grad_outputs):
         condition = self.get_retained_inputs()[0]
         xp = cuda.get_array_module(condition.data)
         g, = grad_outputs
-        zeros = xp.zeros(g.shape, dtype=g.dtype)
+        zero = xp.zeros((), dtype=g.dtype)
         ret = []
         if 0 in indexes:
             ret.append(None)
         if 1 in indexes:
-            gx, = Where().apply((condition.data, g, zeros))
-            ret.append(gx)
+            gx, = Where().apply((condition.data, g, zero))
+            ret.append(chainer.functions.sum_to(gx, self.inputs[1].shape))
         if 2 in indexes:
-            gy, = Where().apply((condition.data, zeros, g))
-            ret.append(gy)
+            gy, = Where().apply((condition.data, zero, g))
+            ret.append(chainer.functions.sum_to(gy, self.inputs[2].shape))
         return ret
 
 
diff --git a/chainer/functions/math/maximum.py b/chainer/functions/math/maximum.py
--- a/chainer/functions/math/maximum.py
+++ b/chainer/functions/math/maximum.py
@@ -15,36 +15,42 @@ def check_type_forward(self, in_types):
         type_check.expect(
             in_types[0].dtype.kind == 'f',
             in_types[0].dtype == in_types[1].dtype,
-            in_types[0].shape == in_types[1].shape
         )
+        type_check.expect_broadcast_shapes(
+            in_types[0].shape, in_types[1].shape)
 
     def forward_cpu(self, inputs):
+        # may broadcast
         self.retain_inputs((0, 1))
         x1, x2 = inputs
         y = numpy.maximum(x1, x2)
         return utils.force_array(y),
 
     def forward_gpu(self, inputs):
+        # may broadcast
         self.retain_inputs((0, 1))
         x1, x2 = inputs
         return cuda.cupy.maximum(x1, x2),
 
     def backward(self, indexes, grad_outputs):
         x1, x2 = self.get_retained_inputs()
-        return MaximumGrad(
-            utils.force_array(x1.data >= x2.data)).apply((grad_outputs[0],))
+        return MaximumGrad(x1.data, x2.data).apply((grad_outputs[0],))
 
 
 class MaximumGrad(function_node.FunctionNode):
 
-    def __init__(self, cond):
-        self.cond = cond
+    def __init__(self, x1, x2):
+        self.cond = x1 >= x2
+        self.x1_shape = x1.shape
+        self.x2_shape = x2.shape
 
     def forward_cpu(self, inputs):
         gy, = inputs
-        gx1 = numpy.where(self.cond, gy, gy.dtype.type(0))
-        gx2 = numpy.where(self.cond, gy.dtype.type(0), gy)
-        return utils.force_array(gx1), utils.force_array(gx2)
+        gx1 = utils.force_array(numpy.where(self.cond, gy, gy.dtype.type(0)))
+        gx2 = utils.force_array(numpy.where(self.cond, gy.dtype.type(0), gy))
+        return (
+            utils.sum_to(gx1, self.x1_shape),
+            utils.sum_to(gx2, self.x2_shape))
 
     def forward_gpu(self, inputs):
         gy, = inputs
@@ -55,11 +61,13 @@ def forward_gpu(self, inputs):
             gx2 = cond ? (T)0.0 : gy;
             ''',
             'maximum_bwd1')(self.cond, gy)
-        return gx1, gx2
+        return (
+            utils.sum_to(gx1, self.x1_shape),
+            utils.sum_to(gx2, self.x2_shape))
 
     def backward(self, indexes, grad_outputs):
         return chainer.functions.where(
-            self.cond, grad_outputs[0], grad_outputs[1]),
+            utils.force_array(self.cond), grad_outputs[0], grad_outputs[1]),
 
 
 def maximum(x1, x2):
diff --git a/chainer/functions/math/minimum.py b/chainer/functions/math/minimum.py
--- a/chainer/functions/math/minimum.py
+++ b/chainer/functions/math/minimum.py
@@ -13,10 +13,12 @@ def check_type_forward(self, in_types):
         type_check.expect(
             in_types[0].dtype.kind == 'f',
             in_types[0].dtype == in_types[1].dtype,
-            in_types[0].shape == in_types[1].shape
         )
+        type_check.expect_broadcast_shapes(
+            in_types[0].shape, in_types[1].shape)
 
     def forward(self, inputs):
+        # may broadcast
         self.retain_inputs((0, 1))
         x1, x2 = inputs
         xp = cuda.get_array_module(x1, x2)
@@ -36,9 +38,9 @@ def __init__(self, x1, x2):
     def forward_cpu(self, inputs):
         gy, = inputs
         x1, x2 = self.x1, self.x2
-        gx1 = gy * (x1 <= x2)
-        gx2 = gy * (x1 > x2)
-        return utils.force_array(gx1), utils.force_array(gx2)
+        gx1 = utils.force_array(gy * (x1 <= x2))
+        gx2 = utils.force_array(gy * (x1 > x2))
+        return utils.sum_to(gx1, x1.shape), utils.sum_to(gx2, x2.shape)
 
     def forward_gpu(self, inputs):
         gy, = inputs
@@ -51,7 +53,7 @@ def forward_gpu(self, inputs):
             'T x1, T x2, T gy', 'T gx1',
             'gx1 = (x1 > x2) ? gy : (T)0.0',
             'minimum_bwd2')(x1, x2, gy)
-        return gx1, gx2
+        return utils.sum_to(gx1, x1.shape), utils.sum_to(gx2, x2.shape)
 
     def backward(self, indexes, grad_outputs):
         x1, x2 = self.x1, self.x2


Test patch:
diff --git a/tests/chainer_tests/functions_tests/array_tests/test_where.py b/tests/chainer_tests/functions_tests/array_tests/test_where.py
--- a/tests/chainer_tests/functions_tests/array_tests/test_where.py
+++ b/tests/chainer_tests/functions_tests/array_tests/test_where.py
@@ -12,20 +12,29 @@
 
 
 @testing.parameterize(*testing.product({
-    'shape': [(3, 2, 4)],
+    'shape': [
+        # c, x, y, output
+        ((3, 2, 4),) * 4,
+        ((4,), (3, 1, 1), (2, 1), (3, 2, 4)),
+    ],
     'dtype': [numpy.float16, numpy.float32, numpy.float32],
 }))
 class TestWhere(unittest.TestCase):
 
     def setUp(self):
-        self.c_data = numpy.random.uniform(-1, 1, self.shape) > 0
+        c_shape, x_shape, y_shape, out_shape = self.shape
+        self.c_data = numpy.random.uniform(-1, 1, c_shape) > 0
         self.x_data = \
-            numpy.random.uniform(-1, 1, self.shape).astype(self.dtype)
+            numpy.random.uniform(-1, 1, x_shape).astype(self.dtype)
         self.y_data = \
-            numpy.random.uniform(-1, 1, self.shape).astype(self.dtype)
+            numpy.random.uniform(-1, 1, y_shape).astype(self.dtype)
         self.g_data = \
-            numpy.random.uniform(-1, 1, self.shape).astype(self.dtype)
+            numpy.random.uniform(-1, 1, out_shape).astype(self.dtype)
         self.check_backward_options = {'dtype': numpy.float64}
+        if self.dtype == numpy.float16:
+            self.check_backward_options.update({
+                'atol': 1e-3, 'rtol': 1e-3,
+            })
 
     def check_forward(self, c_data, x_data, y_data):
         c = chainer.Variable(c_data)
@@ -34,13 +43,9 @@ def check_forward(self, c_data, x_data, y_data):
 
         z = functions.where(c, x, y)
 
-        self.assertEqual(x.data.shape, z.data.shape)
-
-        for i in numpy.ndindex(c.data.shape):
-            if c.data[i]:
-                self.assertEqual(x.data[i], z.data[i])
-            else:
-                self.assertEqual(y.data[i], z.data[i])
+        xp = c.xp
+        z_data_expected = xp.where(c_data, x_data, y_data)
+        testing.assert_allclose(z.array, z_data_expected)
 
     def test_forward_cpu(self):
         self.check_forward(self.c_data, self.x_data, self.y_data)
diff --git a/tests/chainer_tests/functions_tests/math_tests/test_maximum.py b/tests/chainer_tests/functions_tests/math_tests/test_maximum.py
--- a/tests/chainer_tests/functions_tests/math_tests/test_maximum.py
+++ b/tests/chainer_tests/functions_tests/math_tests/test_maximum.py
@@ -1,8 +1,8 @@
 import unittest
 
-import chainer
 import numpy
 
+import chainer
 from chainer.backends import cuda
 from chainer import functions
 from chainer import gradient_check
@@ -12,37 +12,46 @@
 
 
 @testing.parameterize(*testing.product({
-    'shape': [(3, 2), ()],
+    'shape': [
+        # x1, x2, y
+        ((3, 2), (3, 2), (3, 2)),
+        ((), (), ()),
+        ((3, 2), (3, 1), (3, 2)),
+        ((2,), (3, 2), (3, 2)),
+    ],
     'dtype': [numpy.float16, numpy.float32, numpy.float64]
 }))
 class TestMaximum(unittest.TestCase):
 
     def setUp(self):
-        shape = self.shape
-        self.gy = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
-        self.ggx1 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
-        self.ggx2 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
+        x1_shape, x2_shape, y_shape = self.shape
+        self.gy = numpy.random.uniform(-1, 1, y_shape).astype(self.dtype)
+        self.ggx1 = numpy.random.uniform(-1, 1, x1_shape).astype(self.dtype)
+        self.ggx2 = numpy.random.uniform(-1, 1, x2_shape).astype(self.dtype)
         self.check_forward_options = {}
         self.check_backward_options = {'dtype': numpy.float64}
         self.check_double_backward_options = {'dtype': numpy.float64}
         if self.dtype == numpy.float16:
-            eps = 2 ** -3
-            self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
-            self.check_backward_options = {
-                'atol': 1e-2, 'rtol': 1e-1, 'dtype': numpy.float64}
-            self.check_double_backward_options = {
-                'atol': 1e-2, 'rtol': 1e-1, 'dtype': numpy.float64}
-        else:
             eps = 1e-2
+            self.check_forward_options.update({'atol': 1e-4, 'rtol': 1e-3})
+            self.check_backward_options.update({
+                'atol': 1e-2, 'rtol': 1e-2})
+            self.check_double_backward_options.update({
+                'atol': 1e-2, 'rtol': 1e-2})
+        else:
+            eps = 1e-3
         self.check_backward_options['eps'] = eps
-
-        self.x1 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
-        self.x2 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
+        self.check_double_backward_options['eps'] = eps
 
         # Avoid close values for stability in numerical gradient.
-        idx = abs(self.x1 - self.x2) < 2 * eps
-        self.x1[idx] = -0.5
-        self.x2[idx] = 0.5
+        retry = 0
+        while True:
+            self.x1 = numpy.random.uniform(-1, 1, x1_shape).astype(self.dtype)
+            self.x2 = numpy.random.uniform(-1, 1, x2_shape).astype(self.dtype)
+            if (abs(self.x1 - self.x2) >= 2 * eps).all():
+                break
+            retry += 1
+            assert retry <= 10, 'Too many retries to generate inputs'
 
         self.y_expected = numpy.maximum(self.x1, self.x2)
 
diff --git a/tests/chainer_tests/functions_tests/math_tests/test_minimum.py b/tests/chainer_tests/functions_tests/math_tests/test_minimum.py
--- a/tests/chainer_tests/functions_tests/math_tests/test_minimum.py
+++ b/tests/chainer_tests/functions_tests/math_tests/test_minimum.py
@@ -1,8 +1,8 @@
 import unittest
 
-import chainer
 import numpy
 
+import chainer
 from chainer.backends import cuda
 from chainer import functions
 from chainer import gradient_check
@@ -12,36 +12,49 @@
 
 
 @testing.parameterize(*testing.product({
-    'shape': [(3, 2), ()],
+    'shape': [
+        # x1, x2, y
+        ((3, 2), (3, 2), (3, 2)),
+        ((), (), ()),
+        ((3, 2), (3, 1), (3, 2)),
+        ((2,), (3, 2), (3, 2)),
+    ],
     'dtype': [numpy.float16, numpy.float32, numpy.float64]
 }))
 class TestMinimum(unittest.TestCase):
 
     def setUp(self):
-        shape = self.shape
-        self.gy = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
-        self.ggx1 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
-        self.ggx2 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
+        x1_shape, x2_shape, y_shape = self.shape
+        self.gy = numpy.random.uniform(-1, 1, y_shape).astype(self.dtype)
+        self.ggx1 = numpy.random.uniform(-1, 1, x1_shape).astype(self.dtype)
+        self.ggx2 = numpy.random.uniform(-1, 1, x2_shape).astype(self.dtype)
         self.check_forward_options = {}
         self.check_backward_options = {'dtype': numpy.float64}
         self.check_double_backward_options = {'dtype': numpy.float64}
         if self.dtype == numpy.float16:
-            eps = 2 ** -3
-            self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
-            self.check_backward_options = {'atol': 1e-2, 'rtol': 1e-1}
-            self.check_double_backward_options = {'atol': 1e-2, 'rtol': 1e-1}
-        else:
             eps = 1e-2
+            self.check_forward_options.update({'atol': 1e-4, 'rtol': 1e-3})
+            self.check_backward_options.update({
+                'atol': 1e-2, 'rtol': 1e-2})
+            self.check_double_backward_options.update({
+                'atol': 1e-2, 'rtol': 1e-2})
+        else:
+            eps = 1e-3
         self.check_backward_options['eps'] = eps
         self.check_double_backward_options['eps'] = eps
 
-        self.x1 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
-        self.x2 = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
+        self.x1 = numpy.random.uniform(-1, 1, x1_shape).astype(self.dtype)
+        self.x2 = numpy.random.uniform(-1, 1, x2_shape).astype(self.dtype)
 
         # Avoid close values for stability in numerical gradient.
-        idx = abs(self.x1 - self.x2) < 2 * eps
-        self.x1[idx] = -0.5
-        self.x2[idx] = 0.5
+        retry = 0
+        while True:
+            self.x1 = numpy.random.uniform(-1, 1, x1_shape).astype(self.dtype)
+            self.x2 = numpy.random.uniform(-1, 1, x2_shape).astype(self.dtype)
+            if (abs(self.x1 - self.x2) >= 2 * eps).all():
+                break
+            retry += 1
+            assert retry <= 10, 'Too many retries to generate inputs'
 
         self.y_expected = numpy.minimum(self.x1, self.x2)