Merge pull request #936 from yebai/patch-1

Adding Mooncake to the AD list.

Author: ChrisRackauckas

Project.toml

@@ -42,6 +42,7 @@ LoggingExtras = "0.4, 1"
 Lux = "1.12.4"
 MLUtils = "0.4"
 ModelingToolkit = "10"
+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"]

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
 ```

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