Merge pull request #1993 from nouiz/crash_inc_sub_grad

theano/tensor/tests/test_inc_subtensor.py

@@ -83,11 +83,11 @@ class Test_inc_subtensor(unittest.TestCase):
                 f(rng_randX(3, 1), rng_randX(1))
                 # These ones should not
                 self.assertRaises(ValueError,
-                        f, rng_randX(3, 1), rng_randX(2))
+                                  f, rng_randX(3, 1), rng_randX(2))
                 self.assertRaises(ValueError,
-                        f, rng_randX(3, 1), rng_randX(3))
+                                  f, rng_randX(3, 1), rng_randX(3))
                 self.assertRaises(ValueError,
-                        f, rng_randX(3, 1), rng_randX(0))
+                                  f, rng_randX(3, 1), rng_randX(0))
 
     def test_simple_3d(self):
         """Increments or sets part of a tensor by a scalar using full slice and
@@ -100,30 +100,42 @@ class Test_inc_subtensor(unittest.TestCase):
         sl2 = slice(sl2_end)
         sl3 = 2
 
-        for do_set in [True, False]:
-            print "Set", do_set
+        val_a = numpy.ones((5, 3, 4))
+        val_inc = 2.3
+        val_sl2_end = 2
 
-            if do_set:
-                resut = tt.set_subtensor(a[sl1, sl3, sl2], increment)
-            else:
-                resut = tt.inc_subtensor(a[sl1, sl3, sl2], increment)
+        for method in [tt.set_subtensor, tt.inc_subtensor]:
+            print "MethodSet", method
 
-            f = theano.function([a, increment, sl2_end], resut)
+            resut = method(a[sl1, sl3, sl2], increment)
 
-            val_a = numpy.ones((5, 3, 4))
-            val_inc = 2.3
-            val_sl2_end = 2
+            f = theano.function([a, increment, sl2_end], resut)
 
             expected_result = numpy.copy(val_a)
             result = f(val_a, val_inc, val_sl2_end)
 
-            if do_set:
+            if method is tt.set_subtensor:
                 expected_result[:, sl3, :val_sl2_end] = val_inc
             else:
                 expected_result[:, sl3, :val_sl2_end] += val_inc
 
             utt.assert_allclose(result, expected_result)
 
+            # Test when we broadcast the result
+            resut = method(a[sl1, sl2], increment)
+
+            f = theano.function([a, increment, sl2_end], resut)
+
+            expected_result = numpy.copy(val_a)
+            result = f(val_a, val_inc, val_sl2_end)
+
+            if method is tt.set_subtensor:
+                expected_result[:, :val_sl2_end] = val_inc
+            else:
+                expected_result[:, :val_sl2_end] += val_inc
+
+            utt.assert_allclose(result, expected_result)
+
     def test_grad_inc_set(self):
         def inc_slice(*s):
             def just_numeric_args(a, b):
@@ -138,19 +150,24 @@ class Test_inc_subtensor(unittest.TestCase):
         for f_slice in [inc_slice, set_slice]:
             # vector
             utt.verify_grad(
-                    f_slice(slice(2, 4, None)),
-                    (numpy.asarray([0, 1, 2, 3, 4, 5.]),
-                        numpy.asarray([9, 9.]), ))
+                f_slice(slice(2, 4, None)),
+                (numpy.asarray([0, 1, 2, 3, 4, 5.]),
+                 numpy.asarray([9, 9.]), ))
 
             # matrix
             utt.verify_grad(
-                    f_slice(slice(1, 2, None), slice(None, None, None)),
-                    (numpy.asarray([[0, 1], [2, 3], [4, 5.]]),
-                        numpy.asarray([[9, 9.]]), ))
+                f_slice(slice(1, 2, None), slice(None, None, None)),
+                (numpy.asarray([[0, 1], [2, 3], [4, 5.]]),
+                 numpy.asarray([[9, 9.]]), ))
 
-            #single element
+            # single element
             utt.verify_grad(
-                    f_slice(2, 1),
-                    (numpy.asarray([[0, 1], [2, 3], [4, 5.]]),
-                        numpy.asarray(9.),))
+                f_slice(2, 1),
+                (numpy.asarray([[0, 1], [2, 3], [4, 5.]]),
+                 numpy.asarray(9.),))
 
+            # broadcast
+            utt.verify_grad(
+                f_slice(2),
+                (numpy.asarray([[0, 1], [2, 3], [4, 5.]]),
+                 numpy.asarray(9.),))