Merge pull request #2171 from MarcCote/cumsum_3D

theano/sandbox/cuda/tests/test_extra_ops.py

@@ -16,9 +16,8 @@ else:
 from theano import tensor as T
 import numpy as np
 import theano
-from theano import config
 from theano.tensor.extra_ops import cumsum, CumsumOp
-
+import itertools
 
 class TestGpuCumsum(theano.tensor.tests.test_extra_ops.TestCumsumOp):
     mode = mode_with_gpu
@@ -45,68 +44,63 @@ class TestGpuCumsum(theano.tensor.tests.test_extra_ops.TestCumsumOp):
     def test_Strides1D(self):
         x = T.fvector('x')
 
-        # Stepped strides
-        f = theano.function([x], cumsum(x[::2]), mode=self.mode)
-        assert [n for n in f.maker.fgraph.toposort()
-                if isinstance(n.op, GpuCumsum)]
-        a = np.random.randint(10, size=(42,)).astype("float32")
-        assert np.allclose(np.cumsum(a[::2]), f(a))
+        for axis in [0, None]:
+            a = np.random.random((42,)).astype("float32")
+            cumsum_function = theano.function([x], cumsum(x, axis=axis), mode=self.mode)
 
-        # Alternative stepped strides
-        f = theano.function([x], cumsum(x), mode=self.mode)
-        assert [n for n in f.maker.fgraph.toposort()
-                if isinstance(n.op, GpuCumsum)]
-        a = np.random.randint(10, size=(42,)).astype("float32")
-        assert np.allclose(np.cumsum(a[::2]), f(a[::2]))
+            slicings = [slice(None, None, None),    # Normal strides
+                        slice(None, None, 2),       # Stepped strides
+                        slice(None, None, -1),      # Negative strides
+                        ]
 
-        # Negative strides
-        f = theano.function([x], cumsum(x[::-1]), mode=self.mode)
+            # Cartesian product of all slicings to test.
+            for slicing in itertools.product(slicings, repeat=x.ndim):
+                f = theano.function([x], cumsum(x[slicing], axis=axis), mode=self.mode)
                 assert [n for n in f.maker.fgraph.toposort()
                         if isinstance(n.op, GpuCumsum)]
-        a = np.random.randint(10, size=(42,)).astype("float32")
-        assert np.allclose(np.cumsum(a[::-1]), f(a))
+                assert np.allclose(np.cumsum(a[slicing], axis=axis), f(a))
+                assert np.allclose(np.cumsum(a[slicing], axis=axis), cumsum_function(a[slicing]))
 
     def test_Strides2D(self):
         x = T.fmatrix('x')
 
-        for shape_axis, axis in zip([0, 1, 0], [0, 1, None]):
+        for axis in [0, 1, None]:
             a = np.random.random((42, 30)).astype("float32")
+            cumsum_function = theano.function([x], cumsum(x, axis=axis), mode=self.mode)
 
-            # Stepped strides along axis=0
-            f = theano.function([x], cumsum(x[::2], axis=axis), mode=self.mode)
-            assert [n for n in f.maker.fgraph.toposort()
-                    if isinstance(n.op, GpuCumsum)]
-            assert np.allclose(np.cumsum(a[::2], axis=axis), f(a))
+            slicings = [slice(None, None, None),    # Normal strides
+                        slice(None, None, 2),       # Stepped strides
+                        slice(None, None, -1),      # Negative strides
+                        ]
 
-            # Stepped strides along axis=1
-            f = theano.function([x], cumsum(x[:, ::2], axis=axis), mode=self.mode)
+            # Cartesian product of all slicings to test.
+            for slicing in itertools.product(slicings, repeat=x.ndim):
+                f = theano.function([x], cumsum(x[slicing], axis=axis), mode=self.mode)
                 assert [n for n in f.maker.fgraph.toposort()
                         if isinstance(n.op, GpuCumsum)]
-            assert np.allclose(np.cumsum(a[:, ::2], axis=axis), f(a))
+                assert np.allclose(np.cumsum(a[slicing], axis=axis), f(a))
+                assert np.allclose(np.cumsum(a[slicing], axis=axis), cumsum_function(a[slicing]))
 
-            # Alternative stepped strides along axis=0
-            f = theano.function([x], cumsum(x), mode=self.mode)
-            assert [n for n in f.maker.fgraph.toposort()
-                    if isinstance(n.op, GpuCumsum)]
-            assert np.allclose(np.cumsum(a[::2]), f(a[::2]))
+    def test_Strides3D(self):
+        x = T.ftensor3('x')
 
-            # Alternative stepped strides along axis=1
-            f = theano.function([x], cumsum(x), mode=self.mode)
-            assert [n for n in f.maker.fgraph.toposort()
-                    if isinstance(n.op, GpuCumsum)]
-            assert np.allclose(np.cumsum(a[:, ::2]), f(a[:, ::2]))
+        for axis in [0, 1, 2, None]:
+            a = np.random.random((42, 30, 25)).astype("float32")
+            cumsum_function = theano.function([x], cumsum(x, axis=axis), mode=self.mode)
 
-            # Negative strides along axis=0
-            f = theano.function([x], cumsum(x[::-1], axis=axis), mode=self.mode)
-            assert [n for n in f.maker.fgraph.toposort()
-                    if isinstance(n.op, GpuCumsum)]
-            assert np.allclose(np.cumsum(a[::-1], axis=axis), f(a))
+            slicings = [slice(None, None, None),    # Normal strides
+                        slice(None, None, 2),       # Stepped strides
+                        slice(None, None, -1),      # Negative strides
+                        ]
 
-            # Negative strides along axis=1
-            f = theano.function([x], cumsum(x[:, ::-1], axis=axis), mode=self.mode)
+            # Cartesian product of all slicings to test.
+            for slicing in itertools.product(slicings, repeat=x.ndim):
+                f = theano.function([x], cumsum(x[slicing], axis=axis), mode=self.mode)
                 assert [n for n in f.maker.fgraph.toposort()
                         if isinstance(n.op, GpuCumsum)]
-            assert np.allclose(np.cumsum(a[:, ::-1], axis=axis), f(a))
+                assert np.allclose(np.cumsum(a[slicing], axis=axis), f(a))
+                assert np.allclose(np.cumsum(a[slicing], axis=axis), cumsum_function(a[slicing]))
+
 
     def test_GpuCumsum1D(self):
         block_max_size = self.max_threads_dim0 * 2
@@ -163,14 +157,63 @@ class TestGpuCumsum(theano.tensor.tests.test_extra_ops.TestCumsumOp):
             assert np.allclose(np.cumsum(a, axis=axis), f(a))
 
             # Use recursive cumsum
-            a_shape = [5, 3]
+            a_shape = [3, 3]
             a_shape[shape_axis] = block_max_size*(block_max_size+1)+2
-            a = np.ones(a_shape, dtype="float32")
+            a = np.random.random(a_shape).astype("float32")
+            a = np.sign(a-0.5).astype("float32")  # Avoid floating point error
             assert np.allclose(np.cumsum(a, axis=axis), f(a))
 
     def test_GpuCumsum3D(self):
-        # Should not use the GPU version.
+        block_max_size = self.max_threads_dim0 * 2
+
         x = T.ftensor3('x')
+        for shape_axis, axis in zip([0, 1, 2, 0], [0, 1, 2, None]):
+            f = theano.function([x], cumsum(x, axis=axis), mode=self.mode)
+            assert [n for n in f.maker.fgraph.toposort()
+                    if isinstance(n.op, GpuCumsum)]
+
+            # Extensive testing for the first 1025 sizes
+            a_shape = [5, 5, 5]
+            a_shape[shape_axis] = 1025
+            a = np.random.rand(*a_shape).astype("float32")
+            slices = [slice(None), slice(None), slice(None)]
+            for i in xrange(a.shape[shape_axis]):
+                slices[shape_axis] = slice(i)
+                fa = f(a[slices])
+                npa = np.cumsum(a[slices], axis=axis)
+                assert np.allclose(npa, fa)
+
+            # Use multiple GPU threadblocks (along accumulation axis)
+            a_shape = [2, 2, 2]
+            a_shape[shape_axis] = block_max_size+2
+            a = np.random.random(a_shape).astype("float32")
+            assert np.allclose(np.cumsum(a, axis=axis), f(a))
+
+            # Use multiple GPU gridblocks (not along accumulation axis)
+            a_shape = [5, 5, 5]
+            a_shape[(shape_axis+1) % 3] = self.max_grid_size1+1
+            a = np.random.random(a_shape).astype("float32")
+            if axis is None:
+                a = np.sign(a-0.5).astype("float32")  # Avoid floating point error
+            assert np.allclose(np.cumsum(a, axis=axis), f(a))
+
+            a_shape = [5, 5, 5]
+            a_shape[(shape_axis+2) % 3] = self.max_grid_size1+1
+            a = np.random.random(a_shape).astype("float32")
+            if axis is None:
+                a = np.sign(a-0.5).astype("float32")  # Avoid floating point error
+            assert np.allclose(np.cumsum(a, axis=axis), f(a))
+
+            # Use recursive cumsum (along accumulation axis)
+            a_shape = [3, 3, 3]
+            a_shape[shape_axis] = block_max_size*(block_max_size+1)+2
+            a = np.random.random(a_shape).astype("float32")
+            a = np.sign(a-0.5).astype("float32")  # Avoid floating point error
+            assert np.allclose(np.cumsum(a, axis=axis), f(a))
+
+    def test_GpuCumsum4D(self):
+        # Should not use the GPU version.
+        x = T.ftensor4('x')
         f = theano.function([x], cumsum(x, axis=1), mode=self.mode)
         assert [n for n in f.maker.fgraph.toposort()
                 if isinstance(n.op, CumsumOp)]

theano/tensor/tests/test_extra_ops.py

@@ -62,7 +62,7 @@ class TestCumsumOp(utt.InferShapeTester):
         utt.verify_grad(self.op, [a])  # Test axis=None
 
         for axis in range(len(a.shape)):
-            utt.verify_grad(self.op_class(axis=axis), [a])
+            utt.verify_grad(self.op_class(axis=axis), [a], eps=4e-4)
 
 
 class TestCumprodOp(utt.InferShapeTester):