Merge pull request #348 from thiboeri/square_diagonal

theano/sparse/sandbox/sp.py

@@ -137,24 +137,46 @@ class Diag(Op):
 
 diag = Diag()
 
+
 class SquareDiagonal(Op):
-    """Return a square sparse (csc) matrix whose diagonal is given by the dense vector argument.
     """
+    Return a square sparse (csc) matrix whose diagonal
+    is given by the dense vector argument.
+    """
+    def __eq__(self, 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:
+            raise TypeError('data argument must be a vector', diag.type)
+
         return gof.Apply(self, [diag],
-                [sparse.SparseType(dtype = diag.dtype,
-                                 format = 'csc')()])
+                [sparse.SparseType(dtype=diag.dtype, format='csc')()])
+
     def perform(self, node, (diag,), (z,)):
         N, = diag.shape
-        indptr = range(N+1)
+        indptr = range(N + 1)
         indices = indptr[0:N]
-        z[0] = scipy_sparse.csc_matrix((diag, indices, indptr), (N,N), copy=True)
+        z[0] = scipy_sparse.csc_matrix((diag, indices, indptr),
+                                       (N, N), copy=True)
 
     def grad(self, input, (gz,)):
         return [diag(gz)]
 
+    def infer_shape(self, nodes, shapes):
+        diag_length = shapes[0][0]
+        return [(diag_length, diag_length)]
+
 square_diagonal = SquareDiagonal()
 
+
 class ColScaleCSC(Op):
     """
     Scale each columns of a sparse matrix by the corresponding element of a dense vector