Corrections and tests for Diag and SquareDiagonal.

theano/sparse/sandbox/sp.py

@@ -9,6 +9,7 @@ U{http://www-users.cs.umn.edu/~saad/software/SPARSKIT/paper.ps}.
 #### COPIED FROM hpu/icml09/sp.py
 
 import numpy
+import scipy
 from scipy import sparse as scipy_sparse
 
 import theano
@@ -30,70 +31,63 @@ def register_specialize(lopt, *tags, **kwargs):
 
 
 class Diag(Op):
+    """Extract the diagonal of a square sparse matrix as a dense
+    vector.
+
+    :param x: A square sparse matrix in csc format.
+
+    :return: A dense vector representing the diagonal elements.
+
+    :note:
+    - The grad implemented is regular, i.e. not structured, since
+      the output is a dense vector.
     """
-    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)])
 
     def perform(self, node, (x,), (z,)):
-        M, N = x.shape
-        if M != N:
-            raise ValueError("DenseDiag argument not square. Shape:", x.shape)
-
-        assert x.format == 'csc'
-
-        data = x.data
-        indices = x.indices
-        indptr = x.indptr
-
-        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]):
-                if indices[i_idx] == j:
-                    diag[j] += data[i_idx]
-        z[0] = diag
+        N, M = x.shape
+        if N != M:
+            raise ValueError('Diag only apply on square matrix')
+        z[0] = x.diagonal()
 
-    def grad(self, (diag,), (gz,)):
+    def grad(self, (x,), (gz,)):
         return [square_diagonal(gz)]
 
     def infer_shape(self, nodes, shapes):
-        matrix_shape = shapes[0]
-        diag_length = matrix_shape[0]
-        return [(diag_length,)]
+        return [(tensor.minimum(*shapes[0]), )]
 
+    def __str__(self):
+        return self.__class__.__name__
 diag = Diag()
 
 
 class SquareDiagonal(Op):
-    """
-    Return a square sparse (csc) matrix whose diagonal
+    """Return a square sparse (csc) matrix whose diagonal
     is given by the dense vector argument.
+
+    :param x: Dense vector for the diagonal.
+
+    :return: A sparse matrix having `x` as diagonal.
+
+    :note:
+    - The grad implemented is regular, i.e. not structured.
     """
+
     def __eq__(self, other):
-        return (type(self) == type(other))
+        return type(self) == type(other)
 
     def __hash__(self):
         return hash(type(self))
 
-    def __str__(self):
-        return "SquareDiagonal"
-
     def make_node(self, diag):
         diag = tensor.as_tensor_variable(diag)
         if diag.type.ndim != 1:
@@ -102,20 +96,25 @@ class SquareDiagonal(Op):
         return gof.Apply(self, [diag],
                 [sparse.SparseType(dtype=diag.dtype, format='csc')()])
 
-    def perform(self, node, (diag,), (z,)):
-        N, = diag.shape
+    def perform(self, node, inputs, (z,)):
+        diag, o_shape = inputs[0], inputs[0].shape * 2
+
+        N = len(diag)
+        data = diag[:N]
+        indices = range(N)
         indptr = range(N + 1)
-        indices = indptr[0:N]
-        z[0] = scipy_sparse.csc_matrix((diag, indices, indptr),
-                                       (N, N), copy=True)
+        tup = (data, indices, indptr)
+
+        z[0] = scipy.sparse.csc_matrix(tup, copy=True)
 
-    def grad(self, input, (gz,)):
+    def grad(self, inputs, (gz,)):
         return [diag(gz)]
 
     def infer_shape(self, nodes, shapes):
-        diag_length = shapes[0][0]
-        return [(diag_length, diag_length)]
+        return [(shapes[0][0], shapes[0][0])]
 
+    def __str__(self):
+        return self.__class__.__name__
 square_diagonal = SquareDiagonal()
 
 

theano/sparse/sandbox/test_sp.py

@@ -18,6 +18,7 @@ from theano.sparse.sandbox import sp
 from theano.sparse.tests.test_basic import random_lil
 from theano.tests import unittest_tools as utt
 from theano.sparse import verify_grad_sparse
+from theano.sparse.tests.test_basic import sparse_random_inputs
 
 
 class TestSP(unittest.TestCase):
@@ -363,30 +364,87 @@ class TestSP(unittest.TestCase):
                     utt.verify_grad(d, [kvals])
 
 
-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=theano.config.floatX)
-        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()
-        sd = sp.square_diagonal(d)
-        f = theano.function([d], sd)
-        n = numpy.zeros((K,K), dtype='int32')
-        for i in range(K):
-            n[i,i] = i
-
-        assert numpy.all(n == f(range(K)).toarray())
+class DiagTester(utt.InferShapeTester):
+    def setUp(self):
+        super(DiagTester, self).setUp()
+        self.op_class = sp.Diag
+        self.op = sp.diag
+
+    def test_op(self):
+        for format in theano.sparse.sparse_formats:
+            variable, data = sparse_random_inputs(format,
+                                                  shape=(10, 10))
+
+            z = self.op(*variable)
+            assert z.type.broadcastable == (False, )
+
+            f = theano.function(variable, z)
+            tested = f(*data)
+            expected = data[0].toarray().diagonal()
+
+            assert numpy.allclose(tested, expected)
+
+    def test_infer_shape(self):
+        for format in theano.sparse.sparse_formats:
+            variable, data = sparse_random_inputs(format,
+                                                  shape=(10, 10))
+            self._compile_and_check(variable,
+                                    [self.op(*variable)],
+                                    data,
+                                    self.op_class)
+
+    def test_grad(self):
+        for format in theano.sparse.sparse_formats:
+            variable, data = sparse_random_inputs(format,
+                                                  shape=(10, 10))
+            verify_grad_sparse(
+                self.op,
+                data,
+                structured=False)
+
+
+class SquareDiagonalTester(utt.InferShapeTester):
+    def setUp(self):
+        super(SquareDiagonalTester, self).setUp()
+        self.op_class = sp.SquareDiagonal
+        self.op = sp.square_diagonal
+
+    def test_op(self):
+        for format in theano.sparse.sparse_formats:
+            for size in range(5, 9):
+                variable = [tensor.vector()]
+                data = [numpy.random.random(size)]
+
+                f = theano.function(variable, self.op(*variable))
+                tested = f(*data).toarray()
+
+                expected = numpy.diag(*data)
+                assert numpy.allclose(tested, expected)
+                assert tested.dtype == expected.dtype
+                assert tested.shape == expected.shape
+
+    def test_infer_shape(self):
+        for format in theano.sparse.sparse_formats:
+            for size in range(5, 9):
+                variable = [tensor.vector()]
+                data = [numpy.random.random(size)]
+
+                self._compile_and_check(variable,
+                                        [self.op(*variable)],
+                                        data,
+                                        self.op_class)
+
+    def test_grad(self):
+        for format in theano.sparse.sparse_formats:
+            for size in range(5, 9):
+                variable = [tensor.vector()]
+                data = [numpy.random.random(size)]
+
+                verify_grad_sparse(
+                    self.op,
+                    data,
+                    structured=False)
+
 
 def test_ensure_sorted_indices():
     x = 2000

theano/sparse/tests/test_basic.py

@@ -102,7 +102,7 @@ def sparse_random_inputs(format, shape, n=1, out_dtype=None, p=0.5):
         if out_dtype in sparse.discrete_dtypes:
             value = numpy.random.randint(20, size=shape).astype(out_dtype)
         else:
-            value = numpy.random.random(shape)
+            value = numpy.random.random(shape).astype(out_dtype)
         return where * value
 
     variable = [getattr(theano.sparse, format + '_matrix')(dtype=out_dtype)