Compute JVP in line searches

Author: devmotion

.github/workflows/CI.yml

@@ -11,8 +11,8 @@ jobs:
     strategy:
       matrix:
         version:
-          - "min"
-          - "lts"
+          # - "min"
+          # - "lts"
           - "1"
         os:
           - ubuntu-latest

Project.toml

@@ -3,6 +3,7 @@ uuid = "429524aa-4258-5aef-a3af-852621145aeb"
 version = "1.14.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -35,7 +36,7 @@ LineSearches = "7.4.0"
 LinearAlgebra = "<0.0.1, 1.6"
 MathOptInterface = "1.17"
 Measurements = "2.14.1"
-NLSolversBase = "7.9.0"
+NLSolversBase = "8"
 NaNMath = "0.3.2, 1"
 OptimTestProblems = "2.0.3"
 PositiveFactorizations = "0.2.2"
@@ -50,7 +51,6 @@ Test = "<0.0.1, 1.6"
 julia = "1.10"
 
 [extras]
-ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
@@ -68,4 +68,8 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Aqua", "Distributions", "ExplicitImports", "JET", "MathOptInterface", "Measurements", "OptimTestProblems", "Random", "RecursiveArrayTools", "StableRNGs", "LineSearches", "NLSolversBase", "PositiveFactorizations", "ReverseDiff", "ADTypes"]
+test = ["Test", "Aqua", "Distributions", "ExplicitImports", "JET", "MathOptInterface", "Measurements", "OptimTestProblems", "Random", "RecursiveArrayTools", "StableRNGs", "LineSearches", "NLSolversBase", "PositiveFactorizations", "ReverseDiff"]
+
+[sources]
+LineSearches = { url = "https://github.com/devmotion/LineSearches.jl.git", rev = "dmw/jvp" }
+NLSolversBase = { url = "https://github.com/devmotion/NLSolversBase.jl.git", rev = "dmw/jvp" }

docs/src/examples/ipnewton_basics.jl

@@ -22,6 +22,7 @@
 # constraint is unbounded from below or above respectively.
 
 using Optim, NLSolversBase #hide
+import ADTypes #hide
 import NLSolversBase: clear! #hide
 
 # # Constrained optimization with `IPNewton`
@@ -78,7 +79,7 @@ using Test                 #src
 @test Optim.converged(res)      #src
 @test Optim.minimum(res) ≈ 0.25 #src
 
-# Like the rest of Optim, you can also use `autodiff=:forward` and just pass in
+# Like the rest of Optim, you can also use `autodiff=ADTypes.AutoForwardDiff()` and just pass in
 # `fun`.
 
 # If we only want to set lower bounds, use `ux = fill(Inf, 2)`

docs/src/examples/maxlikenlm.jl

@@ -21,7 +21,7 @@
 
 using Optim, NLSolversBase
 using LinearAlgebra: diag
-using ForwardDiff
+using ADTypes, ForwardDiff
 
 #md # !!! tip
 #md #     Add Optim with the following command at the Julia command prompt:
@@ -152,7 +152,7 @@ end
 func = TwiceDifferentiable(
     vars -> Log_Likelihood(x, y, vars[1:nvar], vars[nvar+1]),
     ones(nvar + 1);
-    autodiff = :forward,
+    autodiff = AutoForwardDiff(),
 );
 
 # The above statment accepts 4 inputs: the x matrix, the dependent
@@ -163,7 +163,7 @@ func = TwiceDifferentiable(
 # the error variance.
 #
 # The `ones(nvar+1)` are the starting values for the parameters and
-# the `autodiff=:forward` command performs forward mode automatic
+# the `autodiff=AutoForwardDiff()` command performs forward mode automatic
 # differentiation.
 #
 # The actual optimization of the likelihood function is accomplished

docs/src/user/gradientsandhessians.md

@@ -16,9 +16,8 @@ Automatic differentiation techniques are a middle ground between finite differen
 
 Reverse-mode automatic differentiation can be seen as an automatic implementation of the adjoint method mentioned above, and requires a runtime comparable to only one evaluation of ``f``. It is however considerably more complex to implement, requiring to record the execution of the program to then run it backwards, and incurs a larger overhead.
 
-Forward-mode automatic differentiation is supported through the [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) package by providing the `autodiff=:forward` keyword to `optimize`.
-More generic automatic differentiation is supported thanks to [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl), by setting `autodiff` to any compatible backend object from [ADTypes.jl](https://github.com/SciML/ADTypes.jl).
-For instance, the user can choose `autodiff=AutoReverseDiff()`, `autodiff=AutoEnzyme()`, `autodiff=AutoMooncake()` or `autodiff=AutoZygote()` for a reverse-mode gradient computation, which is generally faster than forward mode on large inputs.
+Generic automatic differentiation is supported thanks to [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl), by setting `autodiff` to any compatible backend object from [ADTypes.jl](https://github.com/SciML/ADTypes.jl).
+For instance, the user can choose `autodiff=AutoForwardDiff()` for forward-mode gradient computation or `autodiff=AutoReverseDiff()`, `autodiff=AutoEnzyme()`, `autodiff=AutoMooncake()` or `autodiff=AutoZygote()` for a reverse-mode gradient computation, which is generally faster than forward mode on large inputs.
 Each of these choices requires loading the corresponding package beforehand.
 
 ## Example
@@ -66,14 +65,16 @@ julia> Optim.minimizer(optimize(f, initial_x, BFGS()))
 ```
 Still looks good. Returning to automatic differentiation, let us try both solvers using this
 method.  We enable [forward mode](https://github.com/JuliaDiff/ForwardDiff.jl) automatic
-differentiation by using the `autodiff = :forward` keyword.
+differentiation by using the `autodiff = ADTypes.AutoForwardDiff()` keyword.
 ```jlcon
-julia> Optim.minimizer(optimize(f, initial_x, BFGS(); autodiff = :forward))
+julia> using ADTypes: AutoForwardDiff
+
+julia> Optim.minimizer(optimize(f, initial_x, BFGS(); autodiff = AutoForwardDiff()))
 2-element Array{Float64,1}:
  1.0
  1.0
 
-julia> Optim.minimizer(optimize(f, initial_x, Newton(); autodiff = :forward))
+julia> Optim.minimizer(optimize(f, initial_x, Newton(); autodiff = AutoForwardDiff()))
 2-element Array{Float64,1}:
  1.0
  1.0

docs/src/user/minimization.md

@@ -26,9 +26,10 @@ If we pass `f` alone, Optim will construct an approximate gradient for us using
 ```jl
 optimize(f, x0, LBFGS())
 ```
-For better performance and greater precision, you can pass your own gradient function. If your objective is written in all Julia code with no special calls to external (that is non-Julia) libraries, you can also use automatic differentiation, by using the `autodiff` keyword and setting it to `:forward`:
+For better performance and greater precision, you can pass your own gradient function. If your objective is written in all Julia code with no special calls to external (that is non-Julia) libraries, you can also use automatic differentiation, by using the `autodiff` keyword and setting it to `ADTypes.AutoForwardDiff()`:
 ```julia
-optimize(f, x0, LBFGS(); autodiff = :forward)
+using ADTypes: AutoForwardDiff
+optimize(f, x0, LBFGS(); autodiff = AutoForwardDiff())
 ```
 
 For the Rosenbrock example, the analytical gradient can be shown to be:

ext/OptimMOIExt.jl

@@ -1,7 +1,8 @@
 module OptimMOIExt
 
 using Optim
-using Optim.LinearAlgebra: rmul! 
+using Optim: ADTypes
+using Optim.LinearAlgebra: rmul!
 import MathOptInterface as MOI
 
 function __init__()
@@ -333,7 +334,7 @@ function MOI.optimize!(model::Optimizer{T}) where {T}
             inplace = true,
         )
     else
-        d = Optim.promote_objtype(method, initial_x, :finite, true, f, g!, h!)
+        d = Optim.promote_objtype(method, initial_x, ADTypes.AutoFiniteDiff(; fdtype = Val(:central)), true, f, g!, h!)
         options = Optim.Options(; Optim.default_options(method)..., options...)
         if nl_constrained || has_bounds
             if nl_constrained

src/Manifolds.jl

@@ -20,9 +20,9 @@ project_tangent(M::Manifold, x) = project_tangent!(M, similar(x), x)
 retract(M::Manifold, x) = retract!(M, copy(x))
 
 # Fake objective function implementing a retraction
-mutable struct ManifoldObjective{T<:NLSolversBase.AbstractObjective} <:
+mutable struct ManifoldObjective{M<:Manifold,T<:NLSolversBase.AbstractObjective} <:
                NLSolversBase.AbstractObjective
-    manifold::Manifold
+    manifold::M
     inner_obj::T
 end
 # TODO: is it safe here to call retract! and change x?
@@ -52,6 +52,20 @@ function NLSolversBase.value_gradient!(obj::ManifoldObjective, x)
     return value(obj.inner_obj)
 end
 
+# In general, we have to compute the gradient/Jacobian separately as it has to be projected
+function NLSolversBase.jvp!(obj::ManifoldObjective, x, v)
+    xin = retract(obj.manifold, x)
+    gradient!(obj.inner_obj, xin)
+    project_tangent!(obj.manifold, gradient(obj.inner_obj), xin)
+    return dot(gradient(obj.inner_obj), v)
+end
+function NLSolversBase.value_jvp!(obj::ManifoldObjective, x, v)
+    xin = retract(obj.manifold, x)
+    value_gradient!(obj.inner_obj, xin)
+    project_tangent!(obj.manifold, gradient(obj.inner_obj), xin)
+    return value(obj.inner_obj), dot(gradient(obj.inner_obj), v)
+end
+
 """Flat Euclidean space {R,C}^N, with projections equal to the identity."""
 struct Flat <: Manifold end
 # all the functions below are no-ops, and therefore the generated code
@@ -62,6 +76,10 @@ retract!(M::Flat, x) = x
 project_tangent(M::Flat, g, x) = g
 project_tangent!(M::Flat, g, x) = g
 
+# Optimizations for `Flat` manifold
+NLSolversBase.jvp!(obj::ManifoldObjective{Flat}, x, v) = jvp!(obj.inner_obj, x, v)
+NLSolversBase.value_jvp!(obj::ManifoldObjective{Flat}, x, v) = value_jvp!(obj.inner_obj, x, v)
+
 """Spherical manifold {|x| = 1}."""
 struct Sphere <: Manifold end
 retract!(S::Sphere, x) = (x ./= norm(x))

src/Optim.jl

@@ -16,6 +16,8 @@ documentation online at http://julianlsolvers.github.io/Optim.jl/stable/ .
 """
 module Optim
 
+import ADTypes
+
 using PositiveFactorizations: Positive # for globalization strategy in Newton
 
 using LineSearches: LineSearches # for globalization strategy in Quasi-Newton algs
@@ -41,6 +43,15 @@ using NLSolversBase:
     TwiceDifferentiableConstraints,
     nconstraints,
     nconstraints_x,
+    value,
+    value!,
+    value!!,
+    gradient,
+    gradient!,
+    value_gradient!,
+    value_gradient!!,
+    jvp!,
+    value_jvp!,
     hessian,
     hessian!,
     hessian!!,