Modify test_fill_diagonal_offset.py to cover every possible case

theano/tensor/tests/test_extra_ops.py

@@ -480,10 +480,9 @@ class TestFillDiagonalOffset(utt.InferShapeTester):
         y = tensor.scalar()
         z = tensor.iscalar()
 
-        test_offset = numpy.random.randint(-5,5)
-
         f = function([x, y, z], fill_diagonal_offset(x, y, z))
-        for shp in [(8, 8), (5, 8), (8, 5)]:
+        for test_offset in (-5, -4, -1, 0, 1, 4, 5):
+            for shp in [(8, 8), (5, 8), (8, 5), (5, 5)]:
                 a = numpy.random.rand(*shp).astype(config.floatX)
                 val = numpy.cast[config.floatX](numpy.random.rand())
                 out = f(a, val, test_offset)
@@ -497,7 +496,7 @@ class TestFillDiagonalOffset(utt.InferShapeTester):
                                             a.shape[0]+test_offset )
 
     def test_gradient(self):
-        test_offset = numpy.random.randint(-5,5)
+        for test_offset in (-5, -4, -1, 0, 1, 4, 5):
             # input 'offset' will not be tested
             def fill_diagonal_with_fix_offset( a, val):
                 return fill_diagonal_offset( a, val, test_offset)
@@ -508,13 +507,23 @@ class TestFillDiagonalOffset(utt.InferShapeTester):
             utt.verify_grad(fill_diagonal_with_fix_offset, 
                         [numpy.random.rand(8, 5), numpy.random.rand()],
                             n_tests=1, rng=TestFillDiagonalOffset.rng)
+            utt.verify_grad(fill_diagonal_with_fix_offset, 
+                        [numpy.random.rand(5, 5), numpy.random.rand()],
+                            n_tests=1, rng=TestFillDiagonalOffset.rng)
 
     def test_infer_shape(self):
         x = tensor.dmatrix()
         y = tensor.dscalar()
         z = tensor.iscalar()
+        for test_offset in (-5, -4, -1, 0, 1, 4, 5):
             self._compile_and_check([x, y, z], [self.op(x, y, z)],
                                     [numpy.random.rand(8, 5),
                                      numpy.random.rand(),
-                                 numpy.random.randint(-5,5)],
+                                     test_offset],
                                      self.op_class )
+            self._compile_and_check([x, y, z], [self.op(x, y, z)],
+                                    [numpy.random.rand(5, 8),
+                                     numpy.random.rand(),
+                                     test_offset],
+                                     self.op_class )
+