feat: Add forward mode Mooncake

Author: gdalle

Project.toml

@@ -1,7 +1,7 @@
 name = "ADTypes"
 uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 authors = ["Vaibhav Dixit <[email protected]>, Guillaume Dalle and contributors"]
-version = "1.14.1"
+version = "1.15.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/index.md

@@ -21,6 +21,7 @@ Algorithmic differentiation:
 ```@docs
 AutoForwardDiff
 AutoPolyesterForwardDiff
+AutoMooncakeForward
 ```
 
 Finite differences:

src/ADTypes.jl

@@ -43,6 +43,7 @@ export AutoChainRules,
        AutoGTPSA,
        AutoModelingToolkit,
        AutoMooncake,
+       AutoMooncakeForward,
        AutoPolyesterForwardDiff,
        AutoReverseDiff,
        AutoSymbolics,

src/dense.jl

@@ -294,6 +294,31 @@ end
 
 mode(::AutoMooncake) = ReverseMode()
 
+"""
+    AutoMooncakeForward
+
+Struct used to select the [Mooncake.jl](https://github.com/compintell/Mooncake.jl) backend for automatic differentiation in forward mode.
+
+Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl).
+
+!!! info
+
+    This type was introduced when forward mode became available in Mooncake.jl. It was kept separate from [`AutoMooncake`](@ref) in order to avoid requiring a breaking release of ADTypes.jl.
+
+# Constructors
+
+    AutoMooncakeForward(; config)
+
+# Fields
+
+  - `config`: either `nothing` or an instance of `Mooncake.Config` -- see the docstring of `Mooncake.Config` for more information. `AutoForwardMooncake(; config=nothing)` is equivalent to `AutoForwardMooncake(; config=Mooncake.Config())`, i.e. the default configuration.
+"""
+Base.@kwdef struct AutoMooncakeForward{Tconfig} <: AbstractADType
+    config::Tconfig
+end
+
+mode(::AutoMooncakeForward) = ForwardMode()
+
 """
     AutoPolyesterForwardDiff{chunksize,T}
 

test/dense.jl

@@ -71,7 +71,8 @@ end
     @test ad.absstep === nothing
     @test ad.dir
 
-    ad = AutoFiniteDiff(; fdtype = Val(:central), fdjtype = Val(:forward), relstep = 1e-3, absstep = 1e-4, dir = false)
+    ad = AutoFiniteDiff(; fdtype = Val(:central), fdjtype = Val(:forward),
+        relstep = 1e-3, absstep = 1e-4, dir = false)
     @test ad isa AbstractADType
     @test ad isa AutoFiniteDiff
     @test mode(ad) isa ForwardMode
@@ -126,13 +127,21 @@ end
 end
 
 @testset "AutoMooncake" begin
-    ad = AutoMooncake(; config=nothing)
+    ad = AutoMooncake(; config = nothing)
     @test ad isa AbstractADType
     @test ad isa AutoMooncake
     @test mode(ad) isa ReverseMode
     @test ad.config === nothing
 end
 
+@testset "AutoMooncakeForward" begin
+    ad = AutoMooncakeForward(; config = nothing)
+    @test ad isa AbstractADType
+    @test ad isa AutoMooncakeForward
+    @test mode(ad) isa ForwardMode
+    @test ad.config === nothing
+end
+
 @testset "AutoPolyesterForwardDiff" begin
     ad = AutoPolyesterForwardDiff()
     @test ad isa AbstractADType

test/runtests.jl

@@ -70,6 +70,8 @@ function every_ad_with_options()
         AutoForwardDiff(chunksize = 3, tag = :tag),
         AutoGTPSA(),
         AutoGTPSA(descriptor = Val(:descriptor)),
+        AutoMooncake(; config = :config),
+        AutoMooncakeForward(; config = :config),
         AutoPolyesterForwardDiff(),
         AutoPolyesterForwardDiff(chunksize = 3, tag = :tag),
         AutoReverseDiff(),