Repository: chainer/chainer

Issue:
Fix `F.where` type checking
Related: #6258

The check for `condition` should be done in `__init__`.


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
@@ -14,9 +14,12 @@ def __init__(self, condition):
         self.condition = condition
 
     def check_type_forward(self, in_types):
-        type_check.expect(in_types.size() == 2)
+        type_check._argname(in_types, ('x', 'y'))
         x_type, y_type = in_types
-        condition = self.condition
+        condition = type_check._make_variable_from_array(
+            # allow scalar `condition`
+            chainer.utils.force_array(self.condition),
+            'condition')
 
         type_check.expect(
             condition.dtype == numpy.bool_,
diff --git a/chainer/utils/type_check.py b/chainer/utils/type_check.py
--- a/chainer/utils/type_check.py
+++ b/chainer/utils/type_check.py
@@ -576,6 +576,18 @@ def make_variable(value, name):
         return Variable(value, name)
 
 
+def _make_variable_from_array(array, name):
+    if not isinstance(array, chainer.get_array_types()):
+        raise InvalidType(
+            'isinstance({}, ndarray)'.format(name),
+            'type({}) == {}'.format(name, type(array)),
+        )
+    if in_light_mode():
+        return array
+    else:
+        return Variable(TypeInfo(array.shape, array.dtype), name)
+
+
 class LightMode(object):
 
     def __enter__(self):


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
@@ -1,7 +1,12 @@
+import unittest
+
 import numpy
+import pytest
 
+import chainer
 from chainer import functions
 from chainer import testing
+from chainer.utils import type_check
 
 
 @testing.parameterize(*testing.product({
@@ -10,7 +15,7 @@
         ((3, 2, 4),) * 4,
         ((4,), (3, 1, 1), (2, 1), (3, 2, 4)),
     ],
-    'dtype': [numpy.float16, numpy.float32, numpy.float32],
+    'dtype': [numpy.float16, numpy.float32, numpy.float64],
 }))
 @testing.inject_backend_tests(
     None,
@@ -56,4 +61,32 @@ def forward(self, inputs, devices):
         return z,
 
 
+class TestWhereTypeCheck(unittest.TestCase):
+
+    def check_forward_raises(self, c_data, x_data, y_data):
+        c = chainer.Variable(c_data)
+        x = chainer.Variable(x_data)
+        y = chainer.Variable(y_data)
+        with pytest.raises(type_check.InvalidType):
+            functions.where(c, x, y)
+
+    def test_cond_int(self):
+        c_data = numpy.zeros(3, numpy.int32)
+        x_data = numpy.zeros(3, numpy.float32)
+        y_data = numpy.zeros(3, numpy.float32)
+        self.check_forward_raises(c_data, x_data, y_data)
+
+    def test_xy_precision(self):
+        c_data = numpy.zeros(3, numpy.bool_)
+        x_data = numpy.zeros(3, numpy.float32)
+        y_data = numpy.zeros(3, numpy.float64)
+        self.check_forward_raises(c_data, x_data, y_data)
+
+    def test_shape(self):
+        c_data = numpy.zeros(3, numpy.bool_)
+        x_data = numpy.zeros(1, numpy.float32)
+        y_data = numpy.zeros(2, numpy.float32)
+        self.check_forward_raises(c_data, x_data, y_data)
+
+
 testing.run_module(__name__, __file__)