diff --git a/Project.toml b/Project.toml index 67f39e158..a6a59bf9e 100644 --- a/Project.toml +++ b/Project.toml @@ -42,6 +42,7 @@ LoggingExtras = "0.4, 1" Lux = "1.12.4" MLUtils = "0.4" ModelingToolkit = "9" +Mooncake = "0.4.138" Optim = ">= 1.4.1" OptimizationBase = "2" OptimizationMOI = "0.5" @@ -99,9 +100,10 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" [targets] test = ["Aqua", "BenchmarkTools", "Boltz", "ComponentArrays", "DiffEqFlux", "Enzyme", "FiniteDiff", "Flux", "ForwardDiff", "Ipopt", "IterTools", "Lux", "MLUtils", "ModelingToolkit", "Optim", "OptimizationMOI", "OptimizationOptimJL", "OptimizationOptimisers", "OrdinaryDiffEqTsit5", "Pkg", "Random", "ReverseDiff", "SafeTestsets", "SciMLSensitivity", "SparseArrays", "SparseDiffTools", - "Symbolics", "Test", "Tracker", "Zygote"] + "Symbolics", "Test", "Tracker", "Zygote", "Mooncake"] diff --git a/docs/src/API/ad.md b/docs/src/API/ad.md index c2caa998b..2626180f9 100644 --- a/docs/src/API/ad.md +++ b/docs/src/API/ad.md @@ -9,6 +9,7 @@ The choices for the auto-AD fill-ins with quick descriptions are: - `AutoFiniteDiff()`: Finite differencing, not optimal but always applicable - `AutoModelingToolkit()`: The fastest choice for large scalar optimizations - `AutoEnzyme()`: Highly performant AD choice for type stable and optimized code + - `AutoMooncake()`: Like Zygote and ReverseDiff, but supports GPU and mutating code ## Automatic Differentiation Choice API @@ -22,4 +23,5 @@ OptimizationBase.AutoZygote OptimizationBase.AutoTracker OptimizationBase.AutoModelingToolkit OptimizationBase.AutoEnzyme +OptimizationBase.AutoMooncake ``` diff --git a/test/ADtests.jl b/test/ADtests.jl index fbeed6bcc..90b08477b 100644 --- a/test/ADtests.jl +++ b/test/ADtests.jl @@ -1,5 +1,5 @@ using Optimization, OptimizationOptimJL, OptimizationMOI, Ipopt, Test -using ForwardDiff, Zygote, ReverseDiff, FiniteDiff, Tracker +using ForwardDiff, Zygote, ReverseDiff, FiniteDiff, Tracker, Mooncake using Enzyme, Random x0 = zeros(2) @@ -35,7 +35,7 @@ end @testset "No constraint" begin for adtype in [AutoEnzyme(), AutoForwardDiff(), AutoZygote(), AutoReverseDiff(), AutoFiniteDiff(), AutoModelingToolkit(), AutoSparseForwardDiff(), - AutoSparseReverseDiff(), AutoSparse(AutoZygote()), AutoModelingToolkit(true, true)] + AutoSparseReverseDiff(), AutoSparse(AutoZygote()), AutoModelingToolkit(true, true), AutoMooncake()] optf = OptimizationFunction(rosenbrock, adtype) prob = OptimizationProblem(optf, x0) @@ -46,16 +46,22 @@ end @test sol.retcode == ReturnCode.Success end - sol = solve(prob, Optim.Newton()) - @test 10 * sol.objective < l1 - if adtype != AutoFiniteDiff() - @test sol.retcode == ReturnCode.Success + # `Newton` requires Hession, which Mooncake doesn't support at the moment. + if adtype != AutoMooncake() + sol = solve(prob, Optim.Newton()) + @test 10 * sol.objective < l1 + if adtype != AutoFiniteDiff() + @test sol.retcode == ReturnCode.Success + end end - sol = solve(prob, Optim.KrylovTrustRegion()) - @test 10 * sol.objective < l1 - if adtype != AutoFiniteDiff() - @test sol.retcode == ReturnCode.Success + # Requires Hession, which Mooncake doesn't support at the moment. + if adtype != AutoMooncake() + sol = solve(prob, Optim.KrylovTrustRegion()) + @test 10 * sol.objective < l1 + if adtype != AutoFiniteDiff() + @test sol.retcode == ReturnCode.Success + end end sol = solve(prob, Optimization.LBFGS(), maxiters = 1000) @@ -67,7 +73,7 @@ end @testset "One constraint" begin for adtype in [AutoEnzyme(), AutoForwardDiff(), AutoZygote(), AutoReverseDiff(), AutoFiniteDiff(), AutoModelingToolkit(), AutoSparseForwardDiff(), - AutoSparseReverseDiff(), AutoSparse(AutoZygote()), AutoModelingToolkit(true, true)] + AutoSparseReverseDiff(), AutoSparse(AutoZygote()), AutoModelingToolkit(true, true), AutoMooncake()] cons = (res, x, p) -> (res[1] = x[1]^2 + x[2]^2 - 1.0; return nothing) optf = OptimizationFunction(rosenbrock, adtype, cons = cons) @@ -77,15 +83,18 @@ end sol = solve(prob, Optimization.LBFGS(), maxiters = 1000) @test 10 * sol.objective < l1 - sol = solve(prob, Ipopt.Optimizer(), max_iter = 1000; print_level = 0) - @test 10 * sol.objective < l1 + # Requires Hession, which Mooncake doesn't support at the moment. + if adtype != AutoMooncake() + sol = solve(prob, Ipopt.Optimizer(), max_iter = 1000; print_level = 0) + @test 10 * sol.objective < l1 + end end end @testset "Two constraints" begin for adtype in [AutoForwardDiff(), AutoZygote(), AutoReverseDiff(), AutoFiniteDiff(), AutoModelingToolkit(), AutoSparseForwardDiff(), - AutoSparseReverseDiff(), AutoSparse(AutoZygote()), AutoModelingToolkit(true, true)] + AutoSparseReverseDiff(), AutoSparse(AutoZygote()), AutoModelingToolkit(true, true), AutoMooncake()] function con2_c(res, x, p) res[1] = x[1]^2 + x[2]^2 res[2] = x[2] * sin(x[1]) - x[1] @@ -99,7 +108,10 @@ end sol = solve(prob, Optimization.LBFGS(), maxiters = 1000) @test 10 * sol.objective < l1 - sol = solve(prob, Ipopt.Optimizer(), max_iter = 1000; print_level = 0) - @test 10 * sol.objective < l1 + # Requires Hession, which Mooncake doesn't support at the moment. + if adtype != AutoMooncake() + sol = solve(prob, Ipopt.Optimizer(), max_iter = 1000; print_level = 0) + @test 10 * sol.objective < l1 + end end end