Add doc and trace test for not square matrix.

theano/sandbox/linalg/ops.py

@@ -395,6 +395,7 @@ solve = Solve() # general solve
 #TODO: Optimizations to replace multiplication by matrix inverse with solve() Op (still unwritten)
 
 class ExtractDiag(Op):
+    """ Return the diagonal of matrix """
     def __init__(self, view=False):
         self.view = view
         if self.view:

theano/sandbox/linalg/tests/test_linalg.py

@@ -137,7 +137,8 @@ def test_trace():
     g = trace(x)
     f = theano.function([x], g)
 
-    m = rng.rand(4, 4).astype(config.floatX)
+    for shp in [(2, 3), (3, 2), (3, 3)]:
+        m = rng.rand(*shp).astype(config.floatX)
         v = numpy.trace(m)
         assert v == f(m)