Merge pull request #110 from SciML/gd/forward_mooncake

feat: Add forward mode Mooncake

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "ADTypes"
 uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 authors = ["Vaibhav Dixit <vaibhavyashdixit@gmail.com>, Guillaume Dalle and contributors"]
-version = "1.16.0"
+version = "1.17.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

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

src/dense.jl

@@ -276,10 +276,16 @@ end
 """
     AutoMooncake
 
-Struct used to select the [Mooncake.jl](https://github.com/compintell/Mooncake.jl) backend for automatic differentiation.
+Struct used to select the [Mooncake.jl](https://github.com/compintell/Mooncake.jl) backend for automatic differentiation in reverse mode.
 
 Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl).
 
+!!! info
+
+    When forward mode became available in Mooncake.jl v0.4.147, another struct called [`AutoMooncakeForward`](@ref) was introduced.
+    It was kept separate to avoid a breaking release of ADTypes.jl.
+    [`AutoMooncake`](@ref) remains for reverse mode only.
+
 # Constructors
 
     AutoMooncake(; config=nothing)
@@ -294,6 +300,33 @@ 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 struct was introduced when forward mode became available in Mooncake.jl v0.4.147.
+    It was kept separate from [`AutoMooncake`](@ref) to avoid a breaking release of ADTypes.jl.
+    [`AutoMooncake`](@ref) remains for reverse mode only.
+
+# Constructors
+
+    AutoMooncakeForward(; config=nothing)
+
+# 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 = nothing
+end
+
+mode(::AutoMooncakeForward) = ForwardMode()
+
 """
     AutoPolyesterForwardDiff{chunksize,T}
 

test/dense.jl

@@ -136,6 +136,16 @@ end
     @test ad.config === nothing
 end
 
+@testset "AutoMooncakeForward" begin
+    ad = AutoMooncakeForward(; config = :config)
+    @test ad isa AbstractADType
+    @test ad isa AutoMooncakeForward
+    @test mode(ad) isa ForwardMode
+    @test ad.config === :config
+    ad = AutoMooncakeForward()
+    @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(),