Support autodiff with DifferentiationInterface (#1131)

* Support autodiff with DifferentiationInterface

* Fix versions

* Typo

* Outdated test code

* Structure test sets for easier debugging

Author: gdalle

.github/workflows/CI.yml

@@ -11,7 +11,7 @@ jobs:
     strategy:
       matrix:
         version:
-          - "1.6"
+          - "1.10"
           - "1"
           - "pre"
         os:

Project.toml

@@ -30,7 +30,7 @@ ForwardDiff = "0.10"
 LineSearches = "7.0.1"
 LinearAlgebra = "<0.0.1, 1.6"
 MathOptInterface = "1.17"
-NLSolversBase = "~7.8.0"
+NLSolversBase = "7.9.0"
 NaNMath = "0.3.2, 1"
 OptimTestProblems = "2.0.3"
 Parameters = "0.10, 0.11, 0.12"
@@ -40,9 +40,10 @@ Random = "<0.0.1, 1.6"
 SparseArrays = "<0.0.1, 1.6"
 StatsBase = "0.29, 0.30, 0.31, 0.32, 0.33, 0.34"
 Test = "<0.0.1, 1.6"
-julia = "1.6"
+julia = "1.10"
 
 [extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
@@ -52,8 +53,9 @@ OptimTestProblems = "cec144fc-5a64-5bc6-99fb-dde8f63e154c"
 PositiveFactorizations = "85a6dd25-e78a-55b7-8502-1745935b8125"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Distributions", "MathOptInterface", "Measurements", "OptimTestProblems", "Random", "RecursiveArrayTools", "StableRNGs", "LineSearches", "NLSolversBase", "PositiveFactorizations"]
+test = ["Test", "Distributions", "MathOptInterface", "Measurements", "OptimTestProblems", "Random", "RecursiveArrayTools", "StableRNGs", "LineSearches", "NLSolversBase", "PositiveFactorizations", "ReverseDiff", "ADTypes"]

docs/src/user/gradientsandhessians.md

@@ -16,7 +16,10 @@ 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`. Reverse-mode automatic differentiation is not supported explicitly yet (although you can use it by writing your own `g!` function). There are a number of implementations in Julia, such as [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl).
+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.
+Each of these choices requires loading the corresponding package beforehand.
 
 ## Example
 

src/multivariate/optimize/interface.jl

@@ -58,26 +58,26 @@ fallback_method(d::OnceDifferentiable) = LBFGS()
 fallback_method(d::TwiceDifferentiable) = Newton()
 
 # promote the objective (tuple of callables or an AbstractObjective) according to method requirement
-promote_objtype(method, initial_x, autodiff::Symbol, inplace::Bool, args...) = error("No default objective type for $method and $args.")
+promote_objtype(method, initial_x, autodiff, inplace::Bool, args...) = error("No default objective type for $method and $args.")
 # actual promotions, notice that (args...) captures FirstOrderOptimizer and NonDifferentiable, etc
-promote_objtype(method::ZerothOrderOptimizer, x, autodiff::Symbol, inplace::Bool, args...) = NonDifferentiable(args..., x, real(zero(eltype(x))))
-promote_objtype(method::FirstOrderOptimizer,  x, autodiff::Symbol, inplace::Bool, f) = OnceDifferentiable(f, x, real(zero(eltype(x))); autodiff = autodiff)
-promote_objtype(method::FirstOrderOptimizer,  x, autodiff::Symbol, inplace::Bool, args...) = OnceDifferentiable(args..., x, real(zero(eltype(x))); inplace = inplace)
-promote_objtype(method::FirstOrderOptimizer,  x, autodiff::Symbol, inplace::Bool, f, g, h) = OnceDifferentiable(f, g, x, real(zero(eltype(x))); inplace = inplace)
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, f) = TwiceDifferentiable(f, x, real(zero(eltype(x))); autodiff = autodiff)
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, f::NotInplaceObjective) = TwiceDifferentiable(f, x, real(zero(eltype(x))))
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, f::InplaceObjective) = TwiceDifferentiable(f, x, real(zero(eltype(x))))
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, f::NLSolversBase.InPlaceObjectiveFGHv) = TwiceDifferentiableHV(f, x)
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, f::NLSolversBase.InPlaceObjectiveFG_Hv) = TwiceDifferentiableHV(f, x)
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, f, g) = TwiceDifferentiable(f, g, x, real(zero(eltype(x))); inplace = inplace, autodiff = autodiff)
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, f, g, h) = TwiceDifferentiable(f, g, h, x, real(zero(eltype(x))); inplace = inplace)
+promote_objtype(method::ZerothOrderOptimizer, x, autodiff, inplace::Bool, args...) = NonDifferentiable(args..., x, real(zero(eltype(x))))
+promote_objtype(method::FirstOrderOptimizer,  x, autodiff, inplace::Bool, f) = OnceDifferentiable(f, x, real(zero(eltype(x))); autodiff = autodiff)
+promote_objtype(method::FirstOrderOptimizer,  x, autodiff, inplace::Bool, args...) = OnceDifferentiable(args..., x, real(zero(eltype(x))); inplace = inplace)
+promote_objtype(method::FirstOrderOptimizer,  x, autodiff, inplace::Bool, f, g, h) = OnceDifferentiable(f, g, x, real(zero(eltype(x))); inplace = inplace)
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, f) = TwiceDifferentiable(f, x, real(zero(eltype(x))); autodiff = autodiff)
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, f::NotInplaceObjective) = TwiceDifferentiable(f, x, real(zero(eltype(x))))
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, f::InplaceObjective) = TwiceDifferentiable(f, x, real(zero(eltype(x))))
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, f::NLSolversBase.InPlaceObjectiveFGHv) = TwiceDifferentiableHV(f, x)
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, f::NLSolversBase.InPlaceObjectiveFG_Hv) = TwiceDifferentiableHV(f, x)
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, f, g) = TwiceDifferentiable(f, g, x, real(zero(eltype(x))); inplace = inplace, autodiff = autodiff)
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, f, g, h) = TwiceDifferentiable(f, g, h, x, real(zero(eltype(x))); inplace = inplace)
 # no-op
-promote_objtype(method::ZerothOrderOptimizer, x, autodiff::Symbol, inplace::Bool, nd::NonDifferentiable)  = nd
-promote_objtype(method::ZerothOrderOptimizer, x, autodiff::Symbol, inplace::Bool, od::OnceDifferentiable) = od
-promote_objtype(method::FirstOrderOptimizer,  x, autodiff::Symbol, inplace::Bool, od::OnceDifferentiable) = od
-promote_objtype(method::ZerothOrderOptimizer, x, autodiff::Symbol, inplace::Bool, td::TwiceDifferentiable) = td
-promote_objtype(method::FirstOrderOptimizer,  x, autodiff::Symbol, inplace::Bool, td::TwiceDifferentiable) = td
-promote_objtype(method::SecondOrderOptimizer, x, autodiff::Symbol, inplace::Bool, td::TwiceDifferentiable) = td
+promote_objtype(method::ZerothOrderOptimizer, x, autodiff, inplace::Bool, nd::NonDifferentiable)  = nd
+promote_objtype(method::ZerothOrderOptimizer, x, autodiff, inplace::Bool, od::OnceDifferentiable) = od
+promote_objtype(method::FirstOrderOptimizer,  x, autodiff, inplace::Bool, od::OnceDifferentiable) = od
+promote_objtype(method::ZerothOrderOptimizer, x, autodiff, inplace::Bool, td::TwiceDifferentiable) = td
+promote_objtype(method::FirstOrderOptimizer,  x, autodiff, inplace::Bool, td::TwiceDifferentiable) = td
+promote_objtype(method::SecondOrderOptimizer, x, autodiff, inplace::Bool, td::TwiceDifferentiable) = td
 
 # if no method or options are present
 function optimize(f,         initial_x::AbstractArray; inplace = true, autodiff = :finite, kwargs...)

src/multivariate/solvers/constrained/ipnewton/ipnewton.jl

@@ -7,10 +7,10 @@ end
 
 Base.summary(::IPNewton) = "Interior Point Newton"
 
-promote_objtype(method::IPNewton, x, autodiff::Symbol, inplace::Bool, f::TwiceDifferentiable) = f
-promote_objtype(method::IPNewton, x, autodiff::Symbol, inplace::Bool, f) = TwiceDifferentiable(f, x, real(zero(eltype(x))); autodiff = autodiff)
-promote_objtype(method::IPNewton, x, autodiff::Symbol, inplace::Bool, f, g) = TwiceDifferentiable(f, g, x, real(zero(eltype(x))); inplace = inplace, autodiff = autodiff)
-promote_objtype(method::IPNewton, x, autodiff::Symbol, inplace::Bool, f, g, h) = TwiceDifferentiable(f, g, h, x, real(zero(eltype(x))); inplace = inplace)
+promote_objtype(method::IPNewton, x, autodiff, inplace::Bool, f::TwiceDifferentiable) = f
+promote_objtype(method::IPNewton, x, autodiff, inplace::Bool, f) = TwiceDifferentiable(f, x, real(zero(eltype(x))); autodiff = autodiff)
+promote_objtype(method::IPNewton, x, autodiff, inplace::Bool, f, g) = TwiceDifferentiable(f, g, x, real(zero(eltype(x))); inplace = inplace, autodiff = autodiff)
+promote_objtype(method::IPNewton, x, autodiff, inplace::Bool, f, g, h) = TwiceDifferentiable(f, g, h, x, real(zero(eltype(x))); inplace = inplace)
 
 # TODO: Add support for InitialGuess from LineSearches
 """

test/general/objective_types.jl

@@ -6,38 +6,41 @@
         for T in (OnceDifferentiable, TwiceDifferentiable)
             odad1 = T(x->5.0, rand(1); autodiff = :finite)
             odad2 = T(x->5.0, rand(1); autodiff = :forward)
+            odad3 = T(x->5.0, rand(1); autodiff = AutoReverseDiff())
             Optim.gradient!(odad1, rand(1))
             Optim.gradient!(odad2, rand(1))
-            #    odad3 = T(x->5., rand(1); autodiff = :reverse)
+            Optim.gradient!(odad3, rand(1))
             @test Optim.gradient(odad1) == [0.0]
             @test Optim.gradient(odad2) == [0.0]
-            #    @test odad3.g == [0.0]
+            @test Optim.gradient(odad3) == [0.0]
         end
 
         for a in (1.0, 5.0)
             xa = rand(1)
             odad1 = OnceDifferentiable(x->a*x[1], xa; autodiff = :finite)
             odad2 = OnceDifferentiable(x->a*x[1], xa; autodiff = :forward)
-        #    odad3 = OnceDifferentiable(x->a*x[1], xa; autodiff = :reverse)
+            odad3 = OnceDifferentiable(x->a*x[1], xa; autodiff = AutoReverseDiff())
             Optim.gradient!(odad1, xa)
             Optim.gradient!(odad2, xa)
+            Optim.gradient!(odad3, xa)
             @test Optim.gradient(odad1) ≈ [a]
             @test Optim.gradient(odad2) == [a]
-        #    @test odad3.g == [a]
+            @test Optim.gradient(odad3) == [a]
         end
         for a in (1.0, 5.0)
             xa = rand(1)
             odad1 = OnceDifferentiable(x->a*x[1]^2, xa; autodiff = :finite)
             odad2 = OnceDifferentiable(x->a*x[1]^2, xa; autodiff = :forward)
-        #    odad3 = OnceDifferentiable(x->a*x[1]^2, xa; autodiff = :reverse)
+            odad3 = OnceDifferentiable(x->a*x[1]^2, xa; autodiff = AutoReverseDiff())
             Optim.gradient!(odad1, xa)
             Optim.gradient!(odad2, xa)
-         @test Optim.gradient(odad1) ≈ 2.0*a*xa
+            Optim.gradient!(odad3, xa)
+            @test Optim.gradient(odad1) ≈ 2.0*a*xa
             @test Optim.gradient(odad2) == 2.0*a*xa
-        #    @test odad3.g == 2.0*a*xa
+            @test Optim.gradient(odad3) == 2.0*a*xa
         end
         for dtype in (OnceDifferentiable, TwiceDifferentiable)
-            for autodiff in (:finite, :forward)
+            for autodiff in (:finite, :forward, AutoReverseDiff())
                 differentiable = dtype(x->sum(x), rand(2); autodiff = autodiff)
                 Optim.value(differentiable)
                 Optim.value!(differentiable, rand(2))

test/runtests.jl

@@ -13,6 +13,9 @@ import NLSolversBase: clear!
 import LinearAlgebra: norm, diag, I, Diagonal, dot, eigen, issymmetric, mul!
 import SparseArrays: normalize!, spdiagm
 
+import ReverseDiff
+using ADTypes: AutoReverseDiff
+
 debug_printing = false
 test_broken = false
 
@@ -244,40 +247,41 @@ function run_optim_tests_constrained(method; convergence_exceptions = (),
     end
 end
 
-
-@testset "special" begin
-    for my_test in special_tests
-        println(my_test)
-        @time include(my_test)
+@testset verbose=true "Optim.jl" begin
+    @testset "special" begin
+        @testset for my_test in special_tests
+            println(my_test)
+            @time include(my_test)
+        end
     end
-end
-@testset "general" begin
-    for my_test in general_tests
-        println(my_test)
-        @time include(my_test)
+    @testset "general" begin
+        @testset for my_test in general_tests
+            println(my_test)
+            @time include(my_test)
+        end
     end
-end
-@testset "univariate" begin
-    for my_test in univariate_tests
-        println(my_test)
-        @time include(my_test)
+    @testset "univariate" begin
+        @testset for my_test in univariate_tests
+            println(my_test)
+            @time include(my_test)
+        end
     end
-end
-@testset "multivariate" begin
-    for my_test in multivariate_tests
-        println(my_test)
-        @time include(my_test)
+    @testset "multivariate" begin
+        @testset for my_test in multivariate_tests
+            println(my_test)
+            @time include(my_test)
+        end
     end
-end
 
-println("Literate examples")
-@time include("examples.jl")
+    println("Literate examples")
+    @time include("examples.jl")
 
-@testset "show method for options" begin
-    o = Optim.Options()
-    @test occursin(" = ", sprint(show, o))
-end
+    @testset "show method for options" begin
+        o = Optim.Options()
+        @test occursin(" = ", sprint(show, o))
+    end
 
-@testset "MOI wrapper" begin
-    include("MOI_wrapper.jl")
+    @testset "MOI wrapper" begin
+        include("MOI_wrapper.jl")
+    end
 end