Reviewed SpSum Op and completed non-grad-related tests.

theano/sparse/sandbox/sp.py

@@ -26,7 +26,7 @@ class SpSum(Op):
     axis = None
     sparse_grad = False
 
-    def __init__(self, axis, sparse_grad):
+    def __init__(self, axis, sparse_grad=True):
         """
         :param sparse_grad: if True, this instance ignores the gradient on matrix elements which are implicitly 0.
         """
@@ -42,6 +42,7 @@ class SpSum(Op):
         #therefore, we *do* want Sums with different sparse_grad values to be merged by the merge optimization.
         # This requires them to compare equal.
         return type(self) == type(other) and self.axis == other.axis
+
     def __hash__(self):
         return 76324 ^ hash(type(self)) ^ hash(self.axis)
 
@@ -53,38 +54,48 @@ class SpSum(Op):
         # At least for small matrices (5x5), the .sum() method of a csc matrix returns a dense matrix
         # as the result whether axis is 0 or 1... weird!
         ###
+        #make sure a sparse type has been given in input
+        assert isinstance(x.type, theano.sparse.SparseType)
+        b=()
+        if self.axis is not None:
+            b=(False,)
+        z = tensor.tensor(broadcastable=b, dtype=x.dtype)
+        return gof.Apply(self, [x], [z])
+
+    def infer_shape(self, node, shapes):
+        r=None
         if self.axis is None:
-            z = tensor.tensor(broadcastable=(), dtype=x.dtype)
+            r=[()]
         elif self.axis == 0:
-            if x.format == 'csc':
-                z = tensor.tensor(broadcastable=(False,), dtype=x.dtype)
-            elif x.format == 'csr':
-                #return SparseVector() #WRITEME!
-                raise NotImplementedError()
-            else:
-                raise NotImplementedError()
-        elif self.axis == 1:
-            if x.format == 'csc':
-                #return SparseVector() #WRITEME!
-                raise NotImplementedError()
-            elif x.format == 'csr':
-                z = tensor.tensor(broadcastable=(False,), dtype=x.dtype)
-            else:
-                raise NotImplementedError()
+            r=[(shapes[0][1],)]
         else:
-            assert False #axis should have been verified by self.__init__
-
-        return gof.Apply(self, [x], [z])
+            r=[(shapes[0][0],)]
+        return r
 
     def perform(self,node, (x,), (z,)):
         if self.axis is None:
             z[0] = numpy.asarray(x.sum())
-        if self.axis == 0:
-            z[0] = numpy.asarray(x.sum(axis=0)).reshape((x.shape[1],))
-        if self.axis == 1:
-            z[0] = numpy.asarray(x.sum(axis=1)).reshape((x.shape[0],))
+        else:
+            s = set(xrange(len(x.shape)))-set([self.axis])
+            myreshape = map((lambda i: x.shape[i]), s)
+            if x.format not in ('csc', 'csr'):
+                x = x.asformat(x.format)
+            z[0] = numpy.asarray(x.sum(axis = self.axis)).reshape(myreshape)
+
+            #case by case code for reference
+            #if self.axis == 0:
+            #    if x.format == 'csc':
+            #        z[0] = numpy.asarray(x.sum(axis=self.axis)).reshape((x.shape[1],))
+            #    else:
+            #        z[0] = numpy.asarray(x.asformat(x.format).sum(axis=self.axis)).reshape((x.shape[1],))
+            #if self.axis == 1:
+            #    if x.format == 'csr':
+            #        z[0] = numpy.asarray(x.sum(axis=self.axis)).reshape((x.shape[0],))
+            #    else:
+            #        z[0] = numpy.asarray(x.asformat(x.format).sum(axis=self.axis)).reshape((x.shape[0],))
 
     def grad(self,(x,), (gz,)):
+        print 'grad (sparse:%s):'%self.sparse_grad,x,gz
         if self.sparse_grad:
             if self.axis is None:
                 return [gz * theano.sparse.sp_ones_like(x)]
@@ -95,7 +106,19 @@ class SpSum(Op):
             else:
                 assert False
         else:
-            raise NotImplementedError()
+            # if sparse_grad is False, this instance does not ignore the
+            # gradient on matrix elements which are implicitly 0,
+            # and return a dense version of the data.
+            if self.axis is None:
+                return [gz * tensor.ones_like(theano.sparse.dense_from_sparse(x))]
+            elif self.axis == 0:
+                raise NotImplementedError('non sparse grad, axis=0')
+                pass
+            elif self.axis == 1:
+                raise NotImplementedError('non sparse grad, axis=1')
+                pass
+            else:
+                assert False
 
 def sp_sum(x, axis=None, sparse_grad=False):
     return SpSum(axis, sparse_grad)(x)

theano/sparse/sandbox/test_sp.py

@@ -360,43 +360,78 @@ class TestSP(unittest.TestCase):
                     utt.verify_grad(d, [kvals])
 
     def test_sp_sum(self):
+
+        print '\n\n*************************************************'
+        print '           TEST SUM'
+        print '*************************************************'
+
+        from theano.sparse.sandbox.sp import SpSum
+
         # TODO: test both grad.
+        rng = numpy.random.RandomState(42)
+        from theano.sparse.basic import SparseFromDense,DenseFromSparse
         for format,cast in [("csc",scipy.sparse.csc_matrix), ("csr",scipy.sparse.csr_matrix)]:
+
+            print 'format: %(format)s'%locals()
+
             x = theano.sparse.SparseType(format=format,
                                          dtype=theano.config.floatX)()
             x_data = numpy.arange(20).reshape(5,4).astype(theano.config.floatX)
 
+            #print 'x_data:',x_data
+            #print 'x_data.sum():',x_data.sum()
 
             # Sum on all axis
+            print 'sum on all axis...'
             z = theano.sparse.sandbox.sp.sp_sum(x)
             assert z.type.broadcastable==()
             f = theano.function([x], z)
             x_val = cast(x_data)
             out = f(x_val)
-            assert out == x_val.sum()
+            expected = x_val.sum()
+            assert out == expected
 
             # Sum on axis 0
-            try:
-                z = theano.sparse.sandbox.sp.sp_sum(x, axis=0)
-                assert z.type.broadcastable==(False,)
-                f = theano.function([x], z)
-                x_val = cast(x_data)
-                out = f(x_val)
-                assert (out == x_val.sum(axis=0)).all()
-            except NotImplementedError:
-                pass
+            print 'sum on axis 0...'
+            z = theano.sparse.sandbox.sp.sp_sum(x, axis=0)
+            assert z.type.broadcastable==(False,)
+            f = theano.function([x], z)
+            x_val = cast(x_data)
+            out = f(x_val)
+            expected=x_val.sum(axis=0)
+            assert (out == expected).all()
 
             # Sum on axis 1
-            try:
-                z = theano.sparse.sandbox.sp.sp_sum(x, axis=1)
-                assert z.type.broadcastable==(False,)
-                f = theano.function([x], z)
-                x_val = cast(x_data)
-                out = f(x_val)
-                expected = numpy.asarray(x_val.sum(axis=1)).reshape(x_val.shape[0])
-                assert (out == expected).all()
-            except NotImplementedError:
-                pass
+            print 'sum on axis 1...'
+            z = theano.sparse.sandbox.sp.sp_sum(x, axis=1)
+            assert z.type.broadcastable==(False,)
+            f = theano.function([x], z)
+            x_val = cast(x_data)
+            out = f(x_val)
+            expected = numpy.asarray(x_val.sum(axis=1)).reshape(x_val.shape[0])
+            assert (out == expected).all()
+
+            # Sparse gradient on Sum on all axis
+            # suspended until som estuff get fixed =/
+            if False:
+                print 'grad on sum on all axis...'
+                def fun(x):
+                    ## verify_grad does not handle sparse data, so here's some casting as a workaround.
+                    # x is a dense matrix: make it sparse
+                    sparse_var = SparseFromDense(format)(x)
+                    # apply op
+                    dense_sum = theano.sparse.sandbox.sp.SpSum(axis=None,sparse_grad=False)(sparse_var)
+                    return dense_sum
+                    # cast back to dense so that verify_grad can work
+                    dense_sum = theano.sparse.DenseFromSparse()(sparse_sum)
+                    return dense_sum
+                x_val = x_data.copy()
+                print type(x_val)
+                import pdb;pdb.set_trace()
+                tensor.verify_grad(fun, [x_val], rng=rng)
+                #utt.verify_grad(SpSum(axis=None), [x_val])
+                print 'ok'
+
 
 def test_diagonal():
     for K in 1, 5:
@@ -503,7 +538,7 @@ def test_col_scale():
         print >> sys.stderr, "WARNING: skipping gradient test because verify_grad doesn't support sparse arguments"
 
 if __name__ == '__main__':
-    if 0:
+    if 1:
         testcase =  TestSP
         suite = unittest.TestLoader()
         suite = suite.loadTestsFromTestCase(testcase)