Add Mooncake extension; extend AD conformance suite with NamedTuple and cache-reuse coverage

[yebai]

Summary

• New AbstractPPLMooncakeExt: gradient/jacobian via Mooncake with cache reuse,
  scalar/vector dispatch, and NamedTuple inputs (wrapping VectorEvaluator and
  NamedTupleEvaluator). Bumps AbstractPPL to 0.15.
• Shared AbstractPPLTestExt conformance suite grows two reusable groups:
  • Val(:namedtuple) — value + gradient over NamedTuple inputs, plus an
    ErrorCase for structure mismatch.
  • Val(:cache_reuse) — three sequential value_and_{gradient,jacobian}!!
    calls against a single Prepared to catch backend cache corruption.
• Evaluators._ad_output_arity (:scalar / :vector) and
  Evaluators._check_ad_input factor out the duplicated output-arity check
  and input validation that appeared across both AD extensions; compile-time
  T <: Integer elision is preserved.
• Empty-input fast paths now tag arity (Val(:scalar) / Val(:vector)) so the
  gradient_prep === nothing arity check still fires for length-0 inputs.

[github-actions]

AbstractPPL.jl documentation for PR #160 is available at:
https://TuringLang.github.io/AbstractPPL.jl/previews/PR160/

[codecov]

Codecov Report

❌ Patch coverage is 96.42857% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.20%. Comparing base (10a4182) to head (88eea2a).

Files with missing lines	Patch %	Lines
ext/AbstractPPLTestExt.jl	88.57%	4 Missing ⚠️
src/evaluators/Evaluators.jl	92.30%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #160      +/-   ##
==========================================
+ Coverage   87.11%   88.20%   +1.09%     
==========================================
  Files          14       15       +1     
  Lines         784      899     +115     
==========================================
+ Hits          683      793     +110     
- Misses        101      106       +5     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[yebai]

@shravanngoswamii, this should be the final PR of the sequel, which adds a native Mooncake backend.

EDIT: There are some minor changes to DI and Mooncake extensions motivated by the DynamicPPL needs.

[shravanngoswamii]

Other than these comments, PR looks good to me. Happy to approve!

[shravanngoswamii]

I may be missing some Mooncake-specific expectation here, but can we check if this method should accept every AbstractVector?

With a SubArray/view, reverse-mode prepares and runs, but the returned gradient is a Mooncake tangent object instead of a normal vector-like gradient. Forward-mode prepares but then errors at call time, and Jacobian also errorslater because Mooncake wants a dense vector input. See:

julia> using AbstractPPL, ADTypes, Mooncake

julia> xs = view([1.0, 2.0, 3.0], :)
3-element view(::Vector{Float64}, :) with eltype Float64:
 1.0
 2.0
 3.0

julia> p_rev = AbstractPPL.prepare(AutoMooncake(), x -> sum(abs2, x), xs);

julia> val, grad = AbstractPPL.value_and_gradient!!(p_rev, xs);

julia> typeof(grad)
Mooncake.Tangent{@NamedTuple{parent::Vector{Float64}, indices::Mooncake.NoTangent, offset1::Mooncake.NoTangent, stride1::Mooncake.NoTangent}}

julia> p_fwd = AbstractPPL.prepare(AutoMooncakeForward(), x -> sum(abs2, x), xs);

julia> AbstractPPL.value_and_gradient!!(p_fwd, xs)
ERROR: MethodError: no method matching setindex!(::Mooncake.Tangent{@NamedTuple{…}}, ::Float64, ::Int64)
The function `setindex!` exists, but no method is defined for this combination of argument types.
Stacktrace:
  [1] _fcache_gradient_seed_tangent(x::SubArray{…}, slot::Int64, cursor::Base.RefValue{…}, dict::IdDict{…})
    @ Mooncake ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1240
  [2] #_fcache_gradient_seed_tangent##0
    @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1284 [inlined]
  [3] ntuple
    @ ./ntuple.jl:51 [inlined]
  [4] _fcache_gradient_seed_tangent
    @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1283 [inlined]
  [5] _fcache_gradient_seed_tangent
    @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1206 [inlined]
  [6] (::Mooncake.var"#_fcache_gradient_chunked!!##0#_fcache_gradient_chunked!!##1"{Tuple{…}})(lane::Int64)
    @ Mooncake ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1792
  [7] ntuple
    @ ./ntuple.jl:19 [inlined]
  [8] _fcache_gradient_chunked!!(cache::Mooncake.ForwardCache{…}, input_primals::Tuple{…})
    @ Mooncake ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1791
  [9] value_and_gradient!!
    @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:2032 [inlined]
 [10] value_and_gradient!!(p::AbstractPPL.Evaluators.Prepared{…}, x::SubArray{…})
    @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:147
 [11] top-level scope
    @ REPL[7]:1
Some type information was truncated. Use `show(err)` to see complete types.

julia> p_vec = AbstractPPL.prepare(AutoMooncake(), x -> [x[1] * x[2], x[2] + x[3]], xs);

julia> AbstractPPL.value_and_jacobian!!(p_vec, xs)
ERROR: ArgumentError: value_and_jacobian!! only supports dense vector inputs; got SubArray{Float64, 1, Vector{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}}, true}
Stacktrace:
 [1] _validate_jacobian_argument(x::SubArray{Float64, 1, Vector{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}}, true})
   @ Mooncake ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:2411
 [2] value_and_jacobian!!
   @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:2522 [inlined]
 [3] value_and_jacobian!!(p::AbstractPPL.Evaluators.Prepared{…}, x::SubArray{…})
   @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:207
 [4] top-level scope
   @ REPL[9]:1
Some type information was truncated. Use `show(err)` to see complete types.

julia> 

For comparison, the DI path returns a normal Vector{Float64} for the same scalar-gradient case with the view input.

Is this expected from Mooncake for views? If yes, maybe this extension should restrict the Mooncake vector path to dense vectors and throw a clear ArgumentError early. If views are meant to work, then maybe the returned tangent needs to be converted back into the gradient shape expected by AbstractPPL. A small regression test with view([1.0, 2.0, 3.0], :) would help make the intended behavior clear.

[shravanngoswamii]

Can we check if AbstractPPL's NamedTupleEvaluator shape check should run once before calling Mooncake here?

For vector inputs, this extension calls _check_ad_input before entering Mooncake, so wrong user inputs get AbstractPPL's normal errors. But for NamedTuple inputs, Mooncake validates its cache first for some mismatches. That means prepared(values) and value_and_gradient!!(prepared, values) can surface different errors for the same bad input.

julia> using AbstractPPL, ADTypes, Mooncake

julia> p = AbstractPPL.prepare(AutoMooncake(), vs -> vs.x^2 + sum(abs2, vs.y), (x=0.0, y=zeros(2)),);

julia> p((y=[1.0, 2.0], x=3.0))
ERROR: ArgumentError: Expected the same NamedTuple structure that was used to prepare this evaluator.
Stacktrace:
 [1] _assert_namedtuple_shape(e::AbstractPPL.Evaluators.NamedTupleEvaluator{…}, values::@NamedTuple{…})
   @ AbstractPPL.Evaluators ~/Work/vectorly-ai/AbstractPPL.jl/src/evaluators/Evaluators.jl:224
 [2] (::AbstractPPL.Evaluators.NamedTupleEvaluator{…})(values::@NamedTuple{…})
   @ AbstractPPL.Evaluators ~/Work/vectorly-ai/AbstractPPL.jl/src/evaluators/Evaluators.jl:207
 [3] (::AbstractPPL.Evaluators.Prepared{…})(x::@NamedTuple{…})
   @ AbstractPPL.Evaluators ~/Work/vectorly-ai/AbstractPPL.jl/src/evaluators/Evaluators.jl:40
 [4] top-level scope
   @ REPL[3]:1
Some type information was truncated. Use `show(err)` to see complete types.

julia> AbstractPPL.value_and_gradient!!(p, (y=[1.0, 2.0], x=3.0))
ERROR: PreparedCacheSpecError:
│ Cached autodiff call has a type mismatch for `x1`.
│ Expected top-level type: @NamedTuple{x::Float64, y::Vector{Float64}}
│ Got top-level type: @NamedTuple{y::Vector{Float64}, x::Float64}
│ Prepared pullback, gradient, derivative, HVP, and Hessian caches must be reused with the same top-level
│ argument types they were prepared with.
└

Stacktrace:
 [1] _throw_prepared_cache_spec_error(kind::Symbol, i::Int64, expected::Type, got::Type)
   @ Mooncake ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1071
 [2] macro expansion
   @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1089 [inlined]
 [3] _validate_prepared_cache_inputs(specs::Tuple{…}, fx::Tuple{…})
   @ Mooncake ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:1078
 [4] #value_and_gradient!!#369
   @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:856 [inlined]
 [5] value_and_gradient!!
   @ ~/.julia/packages/Mooncake/9Ej1i/src/interface.jl:849 [inlined]
 [6] value_and_gradient!!(p::AbstractPPL.Evaluators.Prepared{…}, values::@NamedTuple{…})
   @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:127
 [7] top-level scope
   @ REPL[4]:1
Some type information was truncated. Use `show(err)` to see complete types.

julia> 

Same thing happens for missing/renamed fields. For a nested array size mismatch, the evaluator check does fire, so the behavior currently depends on what kind of NamedTuple mismatch happened.

Is the Mooncake cache error the expected API here, or should this behave like prepared(values)? If the AbstractPPL error is preferred, maybe call Evaluators._assert_namedtuple_shape(p.evaluator, values) before Mooncake.value_and_gradient!!.

[shravanngoswamii]

Can we add a small validation for the shape of raw_gradient_target before building the Mooncake cache?

Right now malformed inputs can throw MethodError / BoundsError. For the single-context scalar case, it also reaches the cache/construction path before failing because contexts is not a tuple.

julia> using AbstractPPL, ADTypes, Mooncake

julia> raw(x, offset) = -0.5 * (x[1] - offset)^2;

julia> problem = x -> raw(x, 0.1);

julia> AbstractPPL.prepare(AutoMooncake(), problem, [0.3]; raw_gradient_target=raw)
ERROR: MethodError: no method matching iterate(::typeof(raw))
The function `iterate` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  iterate(::Compiler.ForwardToBackedgeIterator, ::Int64)
   @ Base ~/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/share/julia/Compiler/src/typeinfer.jl:584
  iterate(::Compiler.ForwardToBackedgeIterator)
   @ Base ~/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/share/julia/Compiler/src/typeinfer.jl:584
  iterate(::Base.MethodSpecializations)
   @ Base runtime_internals.jl:1662
  ...

Stacktrace:
 [1] indexed_iterate(I::Function, i::Int64)
   @ Base ./tuple.jl:165
 [2] prepare(adtype::AutoMooncake{…}, problem::Function, x::Vector{…}; check_dims::Bool, raw_gradient_target::Function)
   @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:104
 [3] top-level scope
   @ REPL[3]:1
Some type information was truncated. Use `show(err)` to see complete types.

julia> AbstractPPL.prepare(AutoMooncake(), problem, [0.3]; raw_gradient_target=(raw,))
ERROR: BoundsError: attempt to access Tuple{typeof(raw)} at index [2]
Stacktrace:
 [1] indexed_iterate(t::Tuple{typeof(raw)}, i::Int64, state::Int64)
   @ Base ./tuple.jl:162
 [2] prepare(adtype::AutoMooncake{…}, problem::Function, x::Vector{…}; check_dims::Bool, raw_gradient_target::Tuple{…})
   @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:104
 [3] top-level scope
   @ REPL[4]:1
Some type information was truncated. Use `show(err)` to see complete types.

julia> AbstractPPL.prepare(AutoMooncake(), problem, [0.3]; raw_gradient_target=(raw, 0.1))
ERROR: MethodError: no method matching (AbstractPPLMooncakeExt.MooncakeCache{:scalar})(::Mooncake.Cache{…}, ::typeof(raw), ::Float64)
The type `AbstractPPLMooncakeExt.MooncakeCache{:scalar}` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (AbstractPPLMooncakeExt.MooncakeCache{A})(::C, ::F, ::CT) where {A, C, F, CT<:Tuple}
   @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:39
  (AbstractPPLMooncakeExt.MooncakeCache{A})(::C) where {A, C}
   @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:36

Stacktrace:
 [1] prepare(adtype::AutoMooncake{…}, problem::Function, x::Vector{…}; check_dims::Bool, raw_gradient_target::Tuple{…})
   @ AbstractPPLMooncakeExt ~/Work/vectorly-ai/AbstractPPL.jl/ext/AbstractPPLMooncakeExt.jl:106
 [2] top-level scope
   @ REPL[5]:1
Some type information was truncated. Use `show(err)` to see complete types.

julia> 

I think a clear ArgumentError would be easier to understand.