Merge pull request #442 from mesnilgr/sparse_diag

theano/sparse/sandbox/sp.py

@@ -108,6 +108,15 @@ class Diag(Op):
     """
     Extract the diagonal of a square sparse matrix as a dense vector.
     """
+    def __eq__(self, other):
+        return (type(self) == type(other))
+
+    def __hash__(self):
+        return hash(type(self))
+
+    def __str__(self):
+        return "Diag"
+
     def make_node(self, x):
         return gof.Apply(self, [x], [tensor.tensor(broadcastable=(False,), dtype=x.dtype)])
 
@@ -125,6 +134,8 @@ class Diag(Op):
         diag = numpy.zeros(N, x.dtype)
 
         #TODO: try using ndarrays and then prune() on the result
+        # it could be optimized in the case the sparse structure
+        # does not allow index duplication
 
         for j in xrange(0, N):
             for i_idx in xrange(indptr[j], indptr[j+1]):
@@ -135,8 +146,12 @@ class Diag(Op):
     def grad(self, (diag,), (gz,)):
         return [square_diagonal(gz)]
 
-diag = Diag()
+    def infer_shape(self, nodes, shapes):
+        matrix_shape = shapes[0] 
+        diag_length = matrix_shape[0]
+        return [(diag_length,)]
 
+diag = Diag()
 
 class SquareDiagonal(Op):
     """

theano/sparse/sandbox/test_sp.py

@@ -464,7 +464,20 @@ def test_remove0():
         mat.eliminate_zeros()
         assert result.size == target.size, 'Matrices sizes differ. Have zeros been removed ?'
 
-def test_diagonal():
+def test_diag():
+    m = theano.sparse.csc_matrix()
+    d = sp.diag(m)
+    f = theano.function([m], d)
+    f2 = theano.function([m], d.shape)
+    for K in 1, 5:
+        np_matrix = numpy.asarray(numpy.reshape(range(K**2),(K,K)),dtype='float64')
+        diag = numpy.diagonal(np_matrix)
+        sp_matrix = scipy.sparse.csc_matrix(np_matrix)
+         
+        assert numpy.all(diag == f(sp_matrix))
+        assert f2(sp_matrix) == diag.shape
+        
+def test_square_diagonal():
     for K in 1, 5:
 
         d = tensor.ivector()
@@ -500,12 +513,25 @@ def test_ensure_sorted_indices():
         sorted_theano = f(sample)
         assert numpy.all(sorted_theano.todense() == sorted_scipy.todense())
 
-def test_diagonal_grad():
+def test_square_diagonal_grad():
     def d(x):
         return sp.sp_sum(sp.square_diagonal(x), sparse_grad=True)
     utt.verify_grad(d, [[0.0, 0.1, 0.2, 0.3]],
             mode=theano.Mode(linker='py', optimizer='fast_compile'))
 
+def test_diag_grad():
+    def d(x):
+        sp_x = theano.sparse.csc_from_dense(x)
+        diag_x = sp.diag(sp_x)
+        return diag_x.sum()
+
+    diag_mat = numpy.zeros((4,4))
+    for idx in xrange(4):
+        diag_mat[idx, idx] += idx * 0.1
+
+    utt.verify_grad(d, [diag_mat],
+            mode=theano.Mode(linker='py', optimizer='fast_compile'))
+
 def test_row_scale():
     x = theano.sparse.csc_dmatrix()
     s = theano.tensor.dvector()