Working for real-eigenvector case

theano/tensor/slinalg.py

@@ -350,6 +350,9 @@ class Expm(Op):
     def grad(self, (A,), (g_out,)):
         return [ExpmGrad()(A, g_out)]
 
+    def infer_shape(self, node, shapes):
+        return [shapes[0]]
+
 
 class ExpmGrad(Op):
     """Gradient of the matrix exponential of a square array
@@ -360,26 +363,24 @@ class ExpmGrad(Op):
             "Scipy not available. Scipy is needed for the Expm op")
         A = as_tensor_variable(A)
         assert A.ndim == 2
-        print(gw.shape)
         out = theano.tensor.matrix(dtype=A.dtype)
         return Apply(self, [A, gw], [out,])
 
-    def perform(self, node, (A, gw), outputs):
-         w, M = scipy.linalg.eig(A)
-         outputs[0][0] = numpy.zeros_like(A)
+    def infer_shape(self, node, shapes):
+        return [shapes[0]]
+
+    def perform(self, node, (A, gw), (out,)):
+        w, M = scipy.linalg.eig(A)
 
-        # Mi = scipy.linalg.inv(M)
-        # G = Mi.dot(gw).dot(M)
+        G = scipy.linalg.solve(M, gw).dot(M)
 
-        # V = np.zeros_like(x)
-        # for i in range(V.shape[0]):
-        #     for j in range(V.shape[1]):
-        #         if i != j:
-        #             V[i,j] = G[i,j] * (exp(w[i]) - exp(w[j])) / (w[i] - w[j])
-        #         else:
-        #             V[i,j] = G[i,i] * exp(w[i])
+        exp_w = numpy.exp(w)
+        V = numpy.subtract.outer(exp_w, exp_w) / numpy.subtract.outer(w, w)
+        V[numpy.diag_indices_from(V)] = exp_w
+        numpy.multiply(V, G, V)
 
-        # outputs[0] = M.dot(V).dot(Mi)
+        Mi = scipy.linalg.inv(M)
+        out[0] = numpy.real(M.dot(V).dot(Mi))
 
 
 def expm(A):

theano/tensor/tests/test_slinalg.py

@@ -208,11 +208,47 @@ def test_expm():
     numpy.testing.assert_array_almost_equal(val, ref)
 
 
-def test_expm_grad():
+def test_expm_grad_1():
+    # with symmetric matrix (real eigenvectors)
     if not imported_scipy:
         raise SkipTest("Scipy needed for the expm op.")
     rng = numpy.random.RandomState(utt.fetch_seed())
-    A = rng.randn(5, 5).astype(config.floatX)
+    A = rng.randn(3, 3).astype(config.floatX)
+    A = A + A.T
+
+    tensor.verify_grad(expm, [A,], rng=rng)
+
+def test_expm_grad_2():
+    # with non-symmetric matrix
+    if not imported_scipy:
+        raise SkipTest("Scipy needed for the expm op.")
+    rng = numpy.random.RandomState(utt.fetch_seed())
+    A = rng.randn(3, 3).astype(config.floatX)
+    A = A
+
+    tensor.verify_grad(expm, [A,], rng=rng)
+
+
+def test_expm_grad_3():
+    if not imported_scipy:
+        raise SkipTest("Scipy needed for the expm op.")
+    from theano.gradient import grad
+    rng = numpy.random.RandomState(utt.fetch_seed())
+    A = rng.randn(3, 3).astype(config.floatX)
+
+    h = 1e-7
+    def e(i,j):
+        v = numpy.zeros((3, 3))
+        v[i, j] = 1
+        return v
+
+    x = tensor.matrix()
+    grad_expm_f = function([x], grad(expm(x)[0,1], x))
+    expm_f = function([x], expm(x)[0,1])
+
+    g = lambda i, j:  (expm_f(A + h*e(i,j)) - expm_f(A)) / h
+    numgrad = numpy.array([[g(i,j) for i in range(3)] for j in range(3)])
+
+    print(grad_expm_f(A))
+    print(numgrad)
 
-    eps = None
-    yield (lambda: utt.verify_grad(expm, [A], 3, rng, eps=eps))