Add inner preparation behavior for AutoMooncakeForward

yebai · web-flow · commit 37c9b4cbf124 · 2025-12-30T23:36:11.000Z
diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl b/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")
