feat: use native Mooncake HVP in forward-over-reverse

.github/workflows/Test.yml

@@ -102,7 +102,6 @@ jobs:
           - GTPSA
           - HyperHessians
           - Mooncake
-          - Mooncake-old
           - PolyesterForwardDiff
           - ReverseDiff
           - SparsityDetector

DifferentiationInterface/CHANGELOG.md

@@ -5,7 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.20...main)
+## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.21...main)
+
+## [0.7.21](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.20...v0.7.21)
+
+### Added
+
+- Use native Mooncake HVP in forward-over-reverse ([#1042](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/1042))
 
 ## [0.7.20](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.19...v0.7.19)
 

DifferentiationInterface/Project.toml

@@ -1,6 +1,6 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
-version = "0.7.20"
+version = "0.7.21"
 authors = ["Guillaume Dalle", "Adrian Hill"]
 
 [deps]
@@ -74,7 +74,7 @@ GPUArraysCore = "0.2"
 GTPSA = "1.4.0"
 HyperHessians = "0.2"
 LinearAlgebra = "1"
-Mooncake = "0.5.1"
+Mooncake = "0.5.25"
 PolyesterForwardDiff = "0.1.2"
 ReverseDiff = "1.15.1"
 SparseArrays = "1"

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -9,12 +9,14 @@ using Mooncake:
     Dual,
     prepare_derivative_cache,
     prepare_gradient_cache,
+    prepare_hvp_cache,
     prepare_pullback_cache,
     primal,
     tangent,
     tangent_type,
     value_and_derivative!!,
     value_and_gradient!!,
+    value_and_hvp!!,
     value_and_pullback!!,
     zero_dual,
     zero_tangent,
@@ -44,6 +46,7 @@ include("onearg.jl")
 include("twoarg.jl")
 include("forward_onearg.jl")
 include("forward_twoarg.jl")
+include("second_order.jl")
 include("differentiate_with.jl")
 
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/second_order.jl

@@ -0,0 +1,108 @@
+struct MooncakeHVPPrep{SIG, Tcache} <: DI.HVPPrep{SIG}
+    _sig::Val{SIG}
+    cache::Tcache
+end
+
+function DI.prepare_hvp_nokwarg(
+        strict::Val, f::F, backend::AutoMooncakeForwardOverReverse, x, tx::NTuple, contexts::Vararg{DI.Context, C}
+    ) where {F, C}
+    _sig = DI.signature(f, backend, x, tx, contexts...; strict)
+    config = get_config(backend)
+    cache = prepare_hvp_cache(f, x, map(DI.unwrap, contexts)...; config)
+    prep = MooncakeHVPPrep(_sig, cache)
+    return prep
+end
+
+function DI.gradient_and_hvp(
+        f::F,
+        prep::MooncakeHVPPrep{Y},
+        backend::AutoMooncakeForwardOverReverse,
+        x,
+        tx::NTuple{1},
+        contexts::Vararg{DI.Context, C},
+    ) where {F, Y, C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    dx = only(tx)
+    if isempty(contexts)
+        _, new_g, new_dg = value_and_hvp!!(
+            prep.cache, f, (dx,), x
+        )
+    else
+        dall = (dx, map(zero_tangent_unwrap, contexts)...)
+        _, (new_g,), (new_dg,) = value_and_hvp!!(
+            prep.cache, f, dall, x, map(DI.unwrap, contexts)...
+        )
+    end
+    return _copy_output(new_g), (_copy_output(new_dg),)
+end
+
+function DI.gradient_and_hvp(
+        f::F,
+        prep::MooncakeHVPPrep{Y},
+        backend::AutoMooncakeForwardOverReverse,
+        x,
+        tx::NTuple,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, Y, C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    gs_and_tg = map(tx) do dx
+        if isempty(contexts)
+            _, new_g, new_dg = value_and_hvp!!(
+                prep.cache, f, dx, x
+            )
+        else
+            dall = (dx, map(zero_tangent_unwrap, contexts)...)
+            _, (new_g,), (new_dg,) = value_and_hvp!!(
+                prep.cache, f, dall, x, map(DI.unwrap, contexts)...
+            )
+        end
+        _copy_output(new_g), _copy_output(new_dg)
+    end
+    g = first(gs_and_tg[1])
+    tg = map(last, gs_and_tg)
+    return g, tg
+end
+
+function DI.gradient_and_hvp!(
+        f::F,
+        g,
+        tg::NTuple,
+        prep::MooncakeHVPPrep,
+        backend::AutoMooncakeForwardOverReverse,
+        x,
+        tx::NTuple,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    new_g, new_tg = DI.gradient_and_hvp(f, prep, backend, x, tx, contexts...)
+    copyto!(g, new_g)
+    foreach(copyto!, tg, new_tg)
+    return g, tg
+end
+
+function DI.hvp(
+        f::F,
+        prep::MooncakeHVPPrep,
+        backend::AutoMooncakeForwardOverReverse,
+        x,
+        tx::NTuple,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    return DI.gradient_and_hvp(f, prep, backend, x, tx, contexts...)[2]
+end
+
+function DI.hvp!(
+        f::F,
+        tg::NTuple,
+        prep::MooncakeHVPPrep,
+        backend::AutoMooncakeForwardOverReverse,
+        x,
+        tx::NTuple,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, C}
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
+    _, new_tg = DI.gradient_and_hvp(f, prep, backend, x, tx, contexts...)
+    foreach(copyto!, tg, new_tg)
+    return tg
+end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl

@@ -1,5 +1,11 @@
+const AutoMooncakeForwardOverReverse{C1, C2} = DI.SecondOrder{
+    AutoMooncakeForward{C1},
+    AutoMooncake{C2},
+}
+
 get_config(::AnyAutoMooncake{Nothing}) = Config()
 get_config(backend::AnyAutoMooncake{<:Config}) = backend.config
+get_config(backend::AutoMooncakeForwardOverReverse) = get_config(DI.outer(backend)) # TODO: mix?
 
 @inline zero_tangent_unwrap(c::DI.Context) = zero_tangent(DI.unwrap(c))
 @inline first_unwrap(c, dc) = (DI.unwrap(c), dc)

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -77,7 +77,7 @@ end
 # Test second-order differentiation (forward-over-reverse)
 test_differentiation(
     [SecondOrder(AutoMooncakeForward(), AutoMooncake())],
-    nomatrix(default_scenarios());
+    nomatrix(default_scenarios(; include_constantified = true));
     excluded = EXCLUDED,
     logging = LOGGING,
     testset_name = "Second order"