feat: support forward-mode Mooncake [experimental]

DifferentiationInterface/docs/src/explanation/backends.md

@@ -12,7 +12,7 @@ We support the following dense backend choices from [ADTypes.jl](https://github.
 - [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences)
 - [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff)
 - [`AutoGTPSA`](@extref ADTypes.AutoGTPSA)
-- [`AutoMooncake`](@extref ADTypes.AutoMooncake)
+- [`AutoMooncake`](@extref ADTypes.AutoMooncake) and [`AutoMooncakeForward`](@extref ADTypes.AutoMooncake)
 - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff)
 - [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff)
 - [`AutoSymbolics`](@extref ADTypes.AutoSymbolics)
@@ -48,6 +48,7 @@ In practice, many AD backends have custom implementations for high-level operato
     | `AutoForwardDiff`          | ✅    | ❌    | ✅     | ✅      | ✅     | ✅      | ✅     | ✅      |
     | `AutoGTPSA`                | ✅    | ❌    | ❌     | ✅      | ✅     | ✅      | ✅     | ✅      |
     | `AutoMooncake`             | ❌    | ✅    | ❌     | ❌      | ❌     | ❌      | ❌     | ❌      |
+    | `AutoMooncakeForward`      | ✅    | ❌    | ❌     | ❌      | ❌     | ❌      | ❌     | ❌      |
     | `AutoPolyesterForwardDiff` | 🔀    | ❌    | 🔀     | ✅      | ✅     | 🔀      | 🔀     | 🔀      |
     | `AutoReverseDiff`          | ❌    | 🔀    | ❌     | ✅      | ✅     | ✅      | ❌     | ❌      |
     | `AutoSymbolics`            | ✅    | ❌    | ✅     | ✅      | ✅     | ✅      | ✅     | ✅      |
@@ -68,6 +69,7 @@ Moreover, each context type is supported by a specific subset of backends:
 | `AutoForwardDiff`          | ✅                  | ✅               |
 | `AutoGTPSA`                | ✅                  | ❌               |
 | `AutoMooncake`             | ✅                  | ✅               |
+| `AutoMooncakeForward`      | ✅                  | ✅               |
 | `AutoPolyesterForwardDiff` | ✅                  | ✅               |
 | `AutoReverseDiff`          | ✅                  | ❌               |
 | `AutoSymbolics`            | ✅                  | ✅               |

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -1,16 +1,22 @@
 module DifferentiationInterfaceMooncakeExt
 
-using ADTypes: ADTypes, AutoMooncake
+using ADTypes: ADTypes, AutoMooncake, AutoMooncakeForward
 import DifferentiationInterface as DI
 using Mooncake:
     Mooncake,
     CoDual,
     Config,
+    Dual,
+    prepare_derivative_cache,
     prepare_gradient_cache,
     prepare_pullback_cache,
+    primal,
+    tangent,
     tangent_type,
+    value_and_derivative!!,
     value_and_gradient!!,
     value_and_pullback!!,
+    zero_dual,
     zero_tangent,
     rdata_type,
     fdata,
@@ -25,17 +31,17 @@ using Mooncake:
     _copy_output,
     _copy_to_output!!
 
-DI.check_available(::AutoMooncake) = true
+const AnyAutoMooncake{C} = Union{AutoMooncake{C},AutoMooncakeForward{C}}
 
-get_config(::AutoMooncake{Nothing}) = Config()
-get_config(backend::AutoMooncake{<:Config}) = backend.config
+DI.check_available(::AnyAutoMooncake) = true
 
-# tangents need to be copied before returning, otherwise they are still aliased in the cache
-mycopy(x::Union{Number,AbstractArray{<:Number}}) = copy(x)
-mycopy(x) = deepcopy(x)
+get_config(::AnyAutoMooncake{Nothing}) = Config()
+get_config(backend::AnyAutoMooncake{<:Config}) = backend.config
 
 include("onearg.jl")
 include("twoarg.jl")
+include("forward_onearg.jl")
+include("forward_twoarg.jl")
 include("differentiate_with.jl")
 
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_onearg.jl

@@ -0,0 +1,92 @@
+## Pushforward
+
+struct MooncakeOneArgPushforwardPrep{SIG,Tcache,DX} <: DI.PushforwardPrep{SIG}
+    _sig::Val{SIG}
+    cache::Tcache
+    dx_righttype::DX
+end
+
+function DI.prepare_pushforward_nokwarg(
+    strict::Val,
+    f::F,
+    backend::AutoMooncakeForward,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C}
+    _sig = DI.signature(f, backend, x, tx, contexts...; strict)
+    config = get_config(backend)
+    # TODO: silence_debug_messages
+    cache = prepare_derivative_cache(f, x, map(DI.unwrap, contexts)...; config.debug_mode)
+    dx_righttype = zero_tangent(x)
+    prep = MooncakeOneArgPushforwardPrep(_sig, cache, dx_righttype)
+    return prep
+end
+
+function DI.value_and_pushforward(
+    f::F,
+    prep::MooncakeOneArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x::X,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C,X}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    ys_and_ty = map(tx) do dx
+        dx_righttype =
+            dx isa tangent_type(X) ? dx : _copy_to_output!!(prep.dx_righttype, dx)
+        y_dual = value_and_derivative!!(
+            prep.cache,
+            zero_dual(f),
+            Dual(x, dx_righttype),
+            map(zero_dual ∘ DI.unwrap, contexts)...,
+        )
+        y = primal(y_dual)
+        dy = _copy_output(tangent(y_dual))
+        return y, dy
+    end
+    y = first(ys_and_ty[1])
+    ty = last.(ys_and_ty)
+    return y, ty
+end
+
+function DI.pushforward(
+    f::F,
+    prep::MooncakeOneArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    return DI.value_and_pushforward(f, prep, backend, x, tx, contexts...)[2]
+end
+
+function DI.value_and_pushforward!(
+    f::F,
+    ty::NTuple,
+    prep::MooncakeOneArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    y, new_ty = DI.value_and_pushforward(f, prep, backend, x, tx, contexts...)
+    foreach(copyto!, ty, new_ty)
+    return y, ty
+end
+
+function DI.pushforward!(
+    f::F,
+    ty::NTuple,
+    prep::MooncakeOneArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    DI.value_and_pushforward!(f, ty, prep, backend, x, tx, contexts...)
+    return ty
+end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_twoarg.jl

@@ -0,0 +1,112 @@
+## Pushforward
+
+struct MooncakeTwoArgPushforwardPrep{SIG,Tcache,DX,DY} <: DI.PushforwardPrep{SIG}
+    _sig::Val{SIG}
+    cache::Tcache
+    dx_righttype::DX
+    dy_righttype::DY
+end
+
+function DI.prepare_pushforward_nokwarg(
+    strict::Val,
+    f!::F,
+    y,
+    backend::AutoMooncakeForward,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C}
+    _sig = DI.signature(f!, y, backend, x, tx, contexts...; strict)
+    config = get_config(backend)
+    # TODO: silence_debug_messages
+    cache = prepare_derivative_cache(
+        f!, y, x, map(DI.unwrap, contexts)...; config.debug_mode
+    )
+    dx_righttype = zero_tangent(x)
+    dy_righttype = zero_tangent(y)
+    prep = MooncakeTwoArgPushforwardPrep(_sig, cache, dx_righttype, dy_righttype)
+    return prep
+end
+
+function DI.value_and_pushforward(
+    f!::F,
+    y,
+    prep::MooncakeTwoArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x::X,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C,X}
+    DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
+    ty = map(tx) do dx
+        dx_righttype =
+            dx isa tangent_type(X) ? dx : _copy_to_output!!(prep.dx_righttype, dx)
+        y_dual = zero_dual(y)
+        value_and_derivative!!(
+            prep.cache,
+            zero_dual(f!),
+            y_dual,
+            Dual(x, dx_righttype),
+            map(zero_dual ∘ DI.unwrap, contexts)...,
+        )
+        dy = _copy_output(tangent(y_dual))
+        return dy
+    end
+    return y, ty
+end
+
+function DI.pushforward(
+    f!::F,
+    y,
+    prep::MooncakeOneArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C}
+    DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
+    return DI.value_and_pushforward(f!, y, prep, backend, x, tx, contexts...)[2]
+end
+
+function DI.value_and_pushforward!(
+    f!::F,
+    y::Y,
+    ty::NTuple,
+    prep::MooncakeOneArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x::X,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C,X,Y}
+    DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
+    foreach(tx, ty) do dx, dy
+        dx_righttype =
+            dx isa tangent_type(X) ? dx : _copy_to_output!!(prep.dx_righttype, dx)
+        dy_righttype =
+            dy isa tangent_type(Y) ? dy : _copy_to_output!!(prep.dy_righttype, dy)
+        value_and_derivative!!(
+            prep.cache,
+            zero_dual(f!),
+            Dual(y, dy_righttype),
+            Dual(x, dx_righttype),
+            map(zero_dual ∘ DI.unwrap, contexts)...,
+        )
+        dy === dy_righttype || copyto!(dy, dy_righttype)
+    end
+    return y, ty
+end
+
+function DI.pushforward!(
+    f!::F,
+    y,
+    ty::NTuple,
+    prep::MooncakeOneArgPushforwardPrep,
+    backend::AutoMooncakeForward,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C};
+) where {F,C}
+    DI.check_prep(f!, y, ty, prep, backend, x, tx, contexts...)
+    DI.pushforward!(f!, y, ty, prep, backend, x, tx, contexts...)
+    return ty
+end

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -28,6 +28,7 @@ using ADTypes:
     AutoForwardDiff,
     AutoGTPSA,
     AutoMooncake,
+    AutoMooncakeForward,
     AutoPolyesterForwardDiff,
     AutoReverseDiff,
     AutoSymbolics,
@@ -115,6 +116,7 @@ export AutoFiniteDifferences
 export AutoForwardDiff
 export AutoGTPSA
 export AutoMooncake
+export AutoMooncakeForward
 export AutoPolyesterForwardDiff
 export AutoReverseDiff
 export AutoSymbolics

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -10,7 +10,11 @@ check_no_implicit_imports(DifferentiationInterface)
 
 LOGGING = get(ENV, "CI", "false") == "false"
 
-backends = [AutoMooncake(; config=nothing), AutoMooncake(; config=Mooncake.Config())]
+backends = [
+    AutoMooncake(; config=nothing),
+    AutoMooncake(; config=Mooncake.Config()),
+    AutoMooncakeForward(; config=nothing);
+]
 
 for backend in backends
     @test check_available(backend)