Rework autodiff implementation

examples/dominant_eigenvector_pytorch.py

@@ -0,0 +1,61 @@
+import numpy as np
+import numpy.random as rnd
+import numpy.linalg as la
+import torch
+
+from pymanopt import Problem
+from pymanopt.tools import decorators
+from pymanopt.manifolds import Sphere
+from pymanopt.solvers import TrustRegions
+
+
+def dominant_eigenvector(A):
+    """
+    Returns the dominant eigenvector of the symmetric matrix A.
+
+    Note: For the same A, this should yield the same as the dominant invariant
+    subspace example with p = 1.
+    """
+    m, n = A.shape
+    assert m == n, "matrix must be square"
+    assert np.allclose(np.sum(A - A.T), 0), "matrix must be symmetric"
+
+    manifold = Sphere(n)
+    solver = TrustRegions()
+
+    A_ = torch.from_numpy(A)
+
+    @decorators.pytorch
+    def cost(x):
+        return -x.matmul(A_.matmul(x))
+
+    problem = Problem(manifold=manifold, cost=cost)
+    xopt = solver.solve(problem)
+    return xopt.squeeze()
+
+
+if __name__ == "__main__":
+    # Generate random problem data.
+    n = 128
+    A = rnd.randn(n, n)
+    A = 0.5 * (A + A.T)
+
+    # Calculate the actual solution by a conventional eigenvalue decomposition.
+    w, v = la.eig(A)
+    x = v[:, np.argmax(w)]
+
+    # Solve the problem with pymanopt.
+    xopt = dominant_eigenvector(A)
+
+    # Make sure both vectors have the same direction. Both are valid
+    # eigenvectors, of course, but for comparison we need to get rid of the
+    # ambiguity.
+    if np.sign(x[0]) != np.sign(xopt[0]):
+        xopt = -xopt
+
+    # Print information about the solution.
+    print('')
+    print("l2-norm of x: %f" % la.norm(x))
+    print("l2-norm of xopt: %f" % la.norm(xopt))
+    print("solution found: %s" % np.allclose(x, xopt, rtol=1e-3))
+    print("l2-error: %f" % la.norm(x - xopt))

examples/linreg.py

@@ -1,6 +1,7 @@
 import autograd.numpy as np
 
 from pymanopt import Problem
+from pymanopt.tools import decorators
 from pymanopt.solvers import TrustRegions
 from pymanopt.manifolds import Euclidean
 
@@ -11,6 +12,7 @@
     Y = np.random.randint(-5, 5, (1, 200))
 
     # Cost function is the squared error
+    @decorators.autograd
     def cost(w):
         return np.sum(np.sum((Y - np.dot(w.T, X))**2))
 

examples/linreg_multiple_autograd.py

@@ -1,6 +1,7 @@
 import autograd.numpy as np
 
 from pymanopt import Problem
+from pymanopt.tools import decorators
 from pymanopt.manifolds import Euclidean
 from pymanopt.solvers import TrustRegions
 
@@ -10,6 +11,7 @@
     X = np.zeros((200, 3))
     y = np.zeros((200, 3))
 
+    @decorators.autograd
     def cost(w):
         return np.sum((y - np.dot(X, w)) ** 2)
     # A solver that involves the hessian

examples/low_rank_matrix_approximation.py

@@ -1,6 +1,8 @@
-from pymanopt.manifolds import FixedRankEmbedded
 import autograd.numpy as np
+
 from pymanopt import Problem
+from pymanopt.tools import decorators
+from pymanopt.manifolds import FixedRankEmbedded
 from pymanopt.solvers import ConjugateGradient
 
 # Let A be a (5 x 4) matrix to be approximated
@@ -20,6 +22,7 @@
 # (b) Definition of a cost function (here using autograd.numpy)
 #       Note that the cost must be defined in terms of u, s and vt, where
 #       X = u * diag(s) * vt.
+@decorators.autograd
 def cost(usv):
     delta = .5
     u = usv[0]

examples/optlog_example.py

@@ -1,6 +1,7 @@
 import autograd.numpy as np
 
 from pymanopt import Problem
+from pymanopt.tools import decorators
 from pymanopt.solvers import SteepestDescent
 from pymanopt.manifolds import Stiefel
 
@@ -11,6 +12,7 @@
     X = np.diag([3, 2, 1]).dot(np.random.randn(3, 200))
 
     # Cost function is the squared reconstruction error
+    @decorators.autograd
     def cost(w):
         return np.sum(np.sum((X - np.dot(w, np.dot(w.T, X)))**2))
 

examples/packing_on_the_sphere.py

@@ -1,6 +1,7 @@
 import autograd.numpy as np
 
 from pymanopt import Problem
+from pymanopt.tools import decorators
 from pymanopt.manifolds import Elliptope
 from pymanopt.solvers import ConjugateGradient
 
@@ -9,6 +10,7 @@ def packing_on_the_sphere(n, k, epsilon):
     manifold = Elliptope(n, k)
     solver = ConjugateGradient(mingradnorm=1e-8, maxiter=1e5)
 
+    @decorators.autograd
     def cost(X):
         Y = np.dot(X, X.T)
         # Shift the exponentials by the maximum value to reduce numerical

examples/pca_autograd.py

@@ -1,6 +1,7 @@
 import autograd.numpy as np
 
 from pymanopt import Problem
+from pymanopt.tools import decorators
 from pymanopt.solvers import TrustRegions
 from pymanopt.manifolds import Stiefel
 
@@ -10,6 +11,7 @@
     X = np.diag([3, 2, 1]).dot(np.random.randn(3, 200))
 
     # Cost function is the squared reconstruction error
+    @decorators.autograd
     def cost(w):
         return np.sum(np.sum((X - np.dot(w, np.dot(w.T, X)))**2))
 

examples/regression_offset_autograd.py

@@ -1,6 +1,7 @@
 import autograd.numpy as np
 
 from pymanopt import Problem
+from pymanopt.tools import decorators
 from pymanopt.solvers import TrustRegions
 from pymanopt.manifolds import Euclidean, Product
 
@@ -13,6 +14,7 @@
     # Note, weights is a tuple/list containing both weight vector w and bias b.
     # This is necessary for autograd to calculate the gradient w.r.t. both
     # arguments in one go.
+    @decorators.autograd
     def cost(weights):
         w = weights[0]
         b = weights[1]

pymanopt/core/problem.py

@@ -4,8 +4,7 @@
 """
 from __future__ import print_function
 
-from pymanopt.tools.autodiff import (AutogradBackend, TheanoBackend,
-                                     TensorflowBackend)
+from pymanopt.tools.autodiff import Function
 
 
 class Problem(object):
@@ -50,18 +49,19 @@ class Problem(object):
     """
     def __init__(self, manifold, cost, egrad=None, ehess=None, grad=None,
                  hess=None, arg=None, precon=None, verbosity=2):
-        self.manifold = manifold
-        # We keep a reference to the original cost function in case we want to
-        # call the `prepare` method twice (for instance, after switching from
-        # a first- to second-order method).
         self._cost = None
+        self._egrad = None
+        self._ehess = None
+        self._grad = None
+        self._hess = None
+
+        self.manifold = manifold
         self._original_cost = cost
-        self._egrad = egrad
-        self._ehess = ehess
-        self._grad = grad
-        self._hess = hess
+        self._original_egrad = egrad
+        self._original_ehess = ehess
+        self._original_grad = grad
+        self._original_hess = hess
         self._arg = arg
-        self._backend = None
 
         if precon is None:
             def precon(x, d):
@@ -70,85 +70,56 @@ def precon(x, d):
 
         self.verbosity = verbosity
 
-        self._backends = list(
-            filter(lambda b: b.is_available(), [
-                TheanoBackend(),
-                AutogradBackend(),
-                TensorflowBackend()
-                ]))
-        self._backend = None
-
-    @property
-    def backend(self):
-        if self._backend is None:
-            for backend in self._backends:
-                if backend.is_compatible(self._original_cost, self._arg):
-                    self._backend = backend
-                    break
-            else:
-                backend_names = [str(backend) for backend in self._backends]
-                if self.verbosity >= 1:
-                    print(backend_names)
-                raise ValueError(
-                    "Cannot determine autodiff backend from cost function of "
-                    "type `{:s}`. Available backends are: {:s}".format(
-                        self._original_cost.__class__.__name__,
-                        ", ".join(backend_names)))
-        return self._backend
-
     @property
     def cost(self):
-        if (self._cost is None and callable(self._original_cost) and
-                not AutogradBackend().is_available()):
-            self._cost = self._original_cost
-
-        elif self._cost is None:
-            if self.verbosity >= 1:
-                print("Compiling cost function...")
-            self._cost = self.backend.compile_function(self._original_cost,
-                                                       self._arg)
-
+        if self._cost is None:
+            self._cost = Function(self._original_cost, self._arg)
         return self._cost
 
+    # FIXME: Since _arg is passed to the problem class, we can't have different
+    #        types of functions/call graphs for cost, gradients and Hessians.
+
     @property
     def egrad(self):
         if self._egrad is None:
-            if self.verbosity >= 1:
-                print("Computing gradient of cost function...")
-            egrad = self.backend.compute_gradient(self._original_cost,
-                                                  self._arg)
-            self._egrad = egrad
+            if self._original_egrad is None:
+                self._egrad = self.cost.compute_gradient()
+            else:
+                self._egrad = Function(self._original_egrad, self._arg)
         return self._egrad
 
     @property
     def grad(self):
         if self._grad is None:
-            # Explicit access forces computation/compilation if necessary.
-            egrad = self.egrad
+            if self._original_grad is None:
+                egrad = self.egrad
 
-            def grad(x):
-                return self.manifold.egrad2rgrad(x, egrad(x))
-            self._grad = grad
+                def grad(x):
+                    return self.manifold.egrad2rgrad(x, egrad(x))
+                self._grad = Function(grad)
+            else:
+                self._grad = Function(self._original_grad, self._arg)
         return self._grad
 
     @property
     def ehess(self):
         if self._ehess is None:
-            if self.verbosity >= 1:
-                print("Computing Hessian of cost function...")
-            ehess = self.backend.compute_hessian(self._original_cost,
-                                                 self._arg)
-            self._ehess = ehess
+            if self._original_ehess is None:
+                self._ehess = self.cost.compute_hessian()
+            else:
+                self._ehess = Function(self._original_ehess, self._arg)
         return self._ehess
 
     @property
     def hess(self):
         if self._hess is None:
-            # Explicit access forces computation if necessary.
-            ehess = self.ehess
+            if self._original_hess is None:
+                ehess = self.ehess
 
-            def hess(x, a):
-                return self.manifold.ehess2rhess(
-                    x, self.egrad(x), ehess(x, a), a)
-            self._hess = hess
+                def hess(x, a):
+                    return self.manifold.ehess2rhess(
+                        x, self.egrad(x), ehess(x, a), a)
+                self._hess = Function(hess)
+            else:
+                self._hess = Function(self._original_hess, self._arg)
         return self._hess