Add sparsereversediff with tape compilation

ext/OptimizationSparseDiffExt.jl

@@ -485,14 +485,24 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
         lag_h, f.lag_hess_prototype)
 end
 
+struct OptimizationSparseReverseTag end
+
 function Optimization.instantiate_function(f, x, adtype::AutoSparseReverseDiff,
     p = SciMLBase.NullParameters(),
     num_cons = 0)
     _f = (θ, args...) -> first(f.f(θ, p, args...))
 
     if f.grad === nothing
-        cfg = ReverseDiff.GradientConfig(x)
-        grad = (res, θ, args...) -> ReverseDiff.gradient!(res, x -> _f(x, args...), θ, cfg)
+        if adtype.compile
+            _tape = ReverseDiff.GradientTape(_f, x)
+            tape = ReverseDiff.compile(_tape)
+            grad = function (res, θ, args...)
+                ReverseDiff.gradient!(res, tape, θ)
+            end
+        else
+            cfg = ReverseDiff.GradientConfig(x)
+            grad = (res, θ, args...) -> ReverseDiff.gradient!(res, x -> _f(x, args...), θ, cfg)
+        end
     else
         grad = (G, θ, args...) -> f.grad(G, θ, p, args...)
     end
@@ -502,9 +512,23 @@ function Optimization.instantiate_function(f, x, adtype::AutoSparseReverseDiff,
     if f.hess === nothing
         hess_sparsity = Symbolics.hessian_sparsity(_f, x)
         hess_colors = SparseDiffTools.matrix_colors(tril(hess_sparsity))
-        hess = function (res, θ, args...)
-            res .= SparseDiffTools.forwarddiff_color_jacobian(θ, colorvec = hess_colors, sparsity = hess_sparsity) do θ
-                ReverseDiff.gradient(x -> _f(x, args...), θ)
+        if adtype.compile
+            T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(x))
+            xdual = ForwardDiff.Dual{typeof(T),eltype(x),length(x)}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), length(x))...,))))
+            h_tape = ReverseDiff.GradientTape(_f, xdual)
+            htape = ReverseDiff.compile(h_tape)
+            function g(res1, θ)
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            jaccfg = ForwardColorJacCache(g, x; tag = typeof(T), colorvec = hess_colors, sparsity = hess_sparsity)
+            hess = function (res, θ, args...)
+                SparseDiffTools.forwarddiff_color_jacobian!(res, g, θ, jaccfg)
+            end
+        else
+            hess = function (res, θ, args...)
+                res .= SparseDiffTools.forwarddiff_color_jacobian(θ, colorvec = hess_colors, sparsity = hess_sparsity) do θ
+                    ReverseDiff.gradient(x -> _f(x, args...), θ)
+                end
             end
         end
     else
@@ -553,10 +577,28 @@ function Optimization.instantiate_function(f, x, adtype::AutoSparseReverseDiff,
         fncs = [(x) -> cons_oop(x)[i] for i in 1:num_cons]
         conshess_sparsity = Symbolics.hessian_sparsity.(fncs, Ref(x))
         conshess_colors = SparseDiffTools.matrix_colors.(conshess_sparsity)
-        cons_h = function (res, θ)
-            for i in 1:num_cons
-                res[i] .= SparseDiffTools.forwarddiff_color_jacobian(θ, colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) do θ
-                    ReverseDiff.gradient(fncs[i], θ)
+        if adtype.compile
+            T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(x))
+            xduals = [ForwardDiff.Dual{typeof(T),eltype(x),maximum(conshess_colors[i])}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), maximum(conshess_colors[i]))...,)))) for i in 1:num_cons]
+            consh_tapes = [ReverseDiff.GradientTape(fncs[i], xduals[i]) for i in 1:num_cons] 
+            conshtapes = ReverseDiff.compile.(consh_tapes)
+            function grad_cons(res1, θ, htape)
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            gs = [(res1, x) -> grad_cons(res1, x, conshtapes[i]) for i in 1:num_cons]
+            jaccfgs = [ForwardColorJacCache(gs[i], x; tag = typeof(T), colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) for i in 1:num_cons]
+            println(jaccfgs)
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    SparseDiffTools.forwarddiff_color_jacobian!(res[i], gs[i], θ, jaccfgs[i])
+                end
+            end
+        else
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    res[i] .= SparseDiffTools.forwarddiff_color_jacobian(θ, colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) do θ
+                        ReverseDiff.gradient(fncs[i], θ)
+                    end
                 end
             end
         end
@@ -585,7 +627,16 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
     _f = (θ, args...) -> first(f.f(θ, cache.p, args...))
 
     if f.grad === nothing
-        grad = (res, θ, args...) -> ReverseDiff.gradient!(res, x -> _f(x, args...), θ)
+        if adtype.compile
+            _tape = ReverseDiff.GradientTape(_f, cache.u0)
+            tape = ReverseDiff.compile(_tape)
+            grad = function (res, θ, args...)
+                ReverseDiff.gradient!(res, tape, θ)
+            end
+        else
+            cfg = ReverseDiff.GradientConfig(cache.u0)
+            grad = (res, θ, args...) -> ReverseDiff.gradient!(res, x -> _f(x, args...), θ, cfg)
+        end
     else
         grad = (G, θ, args...) -> f.grad(G, θ, cache.p, args...)
     end
@@ -595,8 +646,24 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
     if f.hess === nothing
         hess_sparsity = Symbolics.hessian_sparsity(_f, cache.u0)
         hess_colors = SparseDiffTools.matrix_colors(tril(hess_sparsity))
-        hess = function (res, θ, args...)
-            SparseDiffTools.forwarddiff_color_jacobian(res, _θ -> ReverseDiff.gradient(x -> _f(x, args...), _θ), θ, sparsity = hess_sparsity, colorvec = hess_colors)
+        if adtype.compile
+            T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(cache.u0))
+            xdual = ForwardDiff.Dual{typeof(T),eltype(cache.u0),length(cache.u0)}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), length(cache.u0))...,))))
+            h_tape = ReverseDiff.GradientTape(_f, xdual)
+            htape = ReverseDiff.compile(h_tape)
+            function g(res1, θ)
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            jaccfg = ForwardColorJacCache(g, cache.u0; tag = typeof(T), colorvec = hess_colors, sparsity = hess_sparsity)
+            hess = function (res, θ, args...)
+                SparseDiffTools.forwarddiff_color_jacobian!(res, g, θ, jaccfg)
+            end
+        else
+            hess = function (res, θ, args...)
+                res .= SparseDiffTools.forwarddiff_color_jacobian(θ, colorvec = hess_colors, sparsity = hess_sparsity) do θ
+                    ReverseDiff.gradient(x -> _f(x, args...), θ)
+                end
+            end
         end
     else
         hess = (H, θ, args...) -> f.hess(H, θ, cache.p, args...)
@@ -644,10 +711,28 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
         fncs = [(x) -> cons_oop(x)[i] for i in 1:num_cons]
         conshess_sparsity = Symbolics.hessian_sparsity.(fncs, Ref(cache.u0))
         conshess_colors = SparseDiffTools.matrix_colors.(conshess_sparsity)
-        cons_h = function (res, θ)
-            for i in 1:num_cons
-                res[i] .= SparseDiffTools.forwarddiff_color_jacobian(θ, colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) do θ
-                    ReverseDiff.gradient(fncs[i], θ)
+        if adtype.compile
+            T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(cache.u0))
+            xduals = [ForwardDiff.Dual{typeof(T),eltype(cache.u0),maximum(conshess_colors[i])}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), maximum(conshess_colors[i]))...,)))) for i in 1:num_cons]
+            consh_tapes = [ReverseDiff.GradientTape(fncs[i], xduals[i]) for i in 1:num_cons] 
+            conshtapes = ReverseDiff.compile.(consh_tapes)
+            function grad_cons(res1, θ, htape)
+                ReverseDiff.gradient!(res1, htape, θ)
+            end
+            gs = [(res1, x) -> grad_cons(res1, x, conshtapes[i]) for i in 1:num_cons]
+            jaccfgs = [ForwardColorJacCache(gs[i], cache.u0; tag = typeof(T), colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) for i in 1:num_cons]
+            println(jaccfgs)
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    SparseDiffTools.forwarddiff_color_jacobian!(res[i], gs[i], θ, jaccfgs[i])
+                end
+            end
+        else
+            cons_h = function (res, θ)
+                for i in 1:num_cons
+                    res[i] .= SparseDiffTools.forwarddiff_color_jacobian(θ, colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) do θ
+                        ReverseDiff.gradient(fncs[i], θ)
+                    end
                 end
             end
         end

test/ADtests.jl

@@ -521,6 +521,28 @@ sol = solve(prob, Optim.KrylovTrustRegion())
 sol = solve(prob, Optimisers.ADAM(0.1), maxiters = 1000)
 @test 10 * sol.objective < l1
 
+optf = OptimizationFunction(rosenbrock, Optimization.AutoSparseReverseDiff(), cons = con2_c)
+optprob = Optimization.instantiate_function(optf, x0, Optimization.AutoSparseReverseDiff(true),
+    nothing, 2)
+G2 = Array{Float64}(undef, 2)
+optprob.grad(G2, x0)
+@test G1≈G2 rtol=1e-4
+H2 = Array{Float64}(undef, 2, 2)
+optprob.hess(H2, x0)
+@test H1≈H2 rtol=1e-4
+res = Array{Float64}(undef, 2)
+optprob.cons(res, x0)
+@test res≈[0.0, 0.0] atol=1e-4
+optprob.cons(res, [1.0, 2.0])
+@test res ≈ [5.0, 0.682941969615793]
+J = Array{Float64}(undef, 2, 2)
+optprob.cons_j(J, [5.0, 3.0])
+@test all(isapprox(J, [10.0 6.0; -0.149013 -0.958924]; rtol = 1e-3))
+H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
+optprob.cons_h(H3, x0)
+@test H3 ≈ [[2.0 0.0; 0.0 2.0], [-0.0 1.0; 1.0 0.0]]
+
+
 optf = OptimizationFunction(rosenbrock, Optimization.AutoSparseReverseDiff(), cons = con2_c)
 optprob = Optimization.instantiate_function(optf, x0, Optimization.AutoSparseReverseDiff(),
     nothing, 2)
@@ -542,6 +564,7 @@ H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
 optprob.cons_h(H3, x0)
 @test H3 ≈ [[2.0 0.0; 0.0 2.0], [-0.0 1.0; 1.0 0.0]]
 
+
 optf = OptimizationFunction(rosenbrock, Optimization.AutoSparseReverseDiff())
 optprob = Optimization.instantiate_function(optf, x0, Optimization.AutoSparseReverseDiff(),
     nothing)