Skip to content

feat: Add forward mode Mooncake #110

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChrisRackauckas merged 5 commits into from
Aug 13, 2025
Merged

feat: Add forward mode Mooncake #110

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1f735b3
feat: Add forward mode Mooncake
gdalle Jun 5, 2025
2f29265
Merge branch 'main' into gd/forward_mooncake
gdalle Jul 1, 2025
a1b0f1a
Merge remote-tracking branch 'origin/main' into gd/forward_mooncake
gdalle Aug 13, 2025
02a0606
Bump
gdalle Aug 13, 2025
97ffd68
Add deets to AutoMooncake
gdalle Aug 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
1 change: 1 addition & 0 deletions docs/src/index.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ Algorithmic differentiation:
```@docs
AutoForwardDiff
AutoPolyesterForwardDiff
AutoMooncakeForward
```

Finite differences:
Expand Down
1 change: 1 addition & 0 deletions src/ADTypes.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ export AutoChainRules,
AutoGTPSA,
AutoModelingToolkit,
AutoMooncake,
AutoMooncakeForward,
Copy link

@yebai yebai Jun 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For completeness, I suggest we consider introducing AutoMooncakeBackward, then alias AutoMooncake = AutoMooncakeBackward.

Copy link
Collaborator Author

@gdalle gdalle Aug 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think having two names for the same thing introduces unneeded confusion. We can remove AutoMooncakeForward at the next breaking release of ADTypes if there ever is one

AutoPolyesterForwardDiff,
AutoReverseDiff,
AutoSymbolics,
Expand Down
35 changes: 34 additions & 1 deletion src/dense.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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}

Expand Down
10 changes: 10 additions & 0 deletions test/dense.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
Loading