From c6bba565736001e9c41bde2ebd812c76fb388b42 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 30 Jun 2025 11:30:22 +0100 Subject: [PATCH 1/6] Update ad.md --- docs/src/API/ad.md | 2 ++ 1 file changed, 2 insertions(+) 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 ``` From 5c97af2f13219ffa08cfb55ac4b4bdf20f149041 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 9 Jul 2025 22:31:50 +0100 Subject: [PATCH 2/6] adding Mooncake to tests --- test/ADtests.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/ADtests.jl b/test/ADtests.jl index fbeed6bcc..af5105f2c 100644 --- a/test/ADtests.jl +++ b/test/ADtests.jl @@ -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) @@ -67,7 +67,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) @@ -85,7 +85,7 @@ 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] From 4c95ec9c3fd7a7fdbcd610202e6e2bdc1db72263 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 14 Jul 2025 15:05:31 +0100 Subject: [PATCH 3/6] Update ADtests.jl --- test/ADtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/ADtests.jl b/test/ADtests.jl index af5105f2c..745a778f6 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) From 311c3ff2157a138fc377571a0a6682292246fd99 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 14 Jul 2025 15:07:20 +0100 Subject: [PATCH 4/6] Update Project.toml --- Project.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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"] From 66418b30e9ff8c7a5ec0682fbf54a1606b339f1e Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 14 Jul 2025 15:52:20 +0100 Subject: [PATCH 5/6] skip `Newton` for Mooncake --- test/ADtests.jl | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/test/ADtests.jl b/test/ADtests.jl index 745a778f6..c8a94ea49 100644 --- a/test/ADtests.jl +++ b/test/ADtests.jl @@ -46,10 +46,13 @@ 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()) From 71b6c9d1ae0b5e47880b67fa66844c5dd6b1770e Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Fri, 25 Jul 2025 16:30:28 +0100 Subject: [PATCH 6/6] Update ADtests.jl --- test/ADtests.jl | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/test/ADtests.jl b/test/ADtests.jl index c8a94ea49..90b08477b 100644 --- a/test/ADtests.jl +++ b/test/ADtests.jl @@ -55,10 +55,13 @@ end 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) @@ -80,8 +83,11 @@ 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 @@ -102,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