Fix inner preparation behavior for Mooncake (#948)

Author: yebai

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -34,6 +34,7 @@ using Mooncake:
 const AnyAutoMooncake{C} = Union{AutoMooncake{C}, AutoMooncakeForward{C}}
 
 DI.check_available(::AnyAutoMooncake{C}) where {C} = true
+DI.inner_preparation_behavior(::AutoMooncakeForward) = DI.PrepareInnerSimple()
 
 include("utils.jl")
 include("onearg.jl")

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -28,6 +28,22 @@ test_differentiation(
     logging = LOGGING,
 );
 
+EXCLUDED = @static if VERSION ≥ v"1.11-" && VERSION ≤ v"1.12-"
+    # testing only :hessian on 1.11 due to an opaque closure bug.
+    # this is potentially the same issue as discussed in
+    # https://github.com/chalk-lab/MistyClosures.jl/pull/12#issue-3278662295
+    [FIRST_ORDER..., :hvp, :second_derivative]
+else
+    [FIRST_ORDER...]
+end
+
+# Test second-order differentiation (forward-over-reverse)
+test_differentiation(
+    [SecondOrder(AutoMooncakeForward(; config = nothing), AutoMooncake(; config = nothing))],
+    excluded = EXCLUDED,
+    logging = true,
+)
+
 @testset "NamedTuples" begin
     ps = (; A = rand(5), B = rand(5))
     myfun(ps) = sum(ps.A .* ps.B)