Add evaluator interface with AD backend extensions

[yebai]

Summary

Adds a prepared evaluator API to AbstractPPL that decouples log-density
evaluation from automatic differentiation, following the interface sketched
in docs/src/interface.md.

Autograd backends are optional. For example, removing the Enzyme package extension would simply fall back to DifferentiationInterface for AutoEnzyme. The broader aim is to rely only on public APIs of autodiff backends, while keeping gradient definitions independent of any particular backend implementation.

• This PR implements this design note.
• Fix rename LogDensityFunction to NamedLogDensity DynamicPPL.jl#880

TODOs

• make sure the evaluator interface is sufficient for DynamicPPL
• make sure the evaluator interface is sufficient for AdvancedVI
• make sure the evaluator interface is sufficient for Bijectors

 using AbstractPPL, ADTypes

 # 1. Define a problem type and its structural evaluator
 struct MyProblem end
 struct MyEvaluator end

 AbstractPPL.prepare(::MyProblem, ::NamedTuple) = MyEvaluator()
 (::MyEvaluator)(nt::NamedTuple) = nt.x^2 + sum(nt.y .^ 2)

 # 2. Prepare an AD-aware evaluator (works with any supported backend)
 using ForwardDiff
 prepared = prepare(AutoForwardDiff(), MyProblem(), (x=0.0, y=[0.0, 0.0]))

 # 3. Evaluate and differentiate — gradients keep the named structure
 val, grad = value_and_gradient(prepared, (x=3.0, y=[1.0, 2.0]))
 # val    = 14.0
 # grad.x = 6.0
 # grad.y = [2.0, 4.0]

 # 4. Verify gradients against finite differences (in tests)
using FiniteDifferences
 test_autograd(prepared, (x=3.0, y=[1.0, 2.0]))

[github-actions]

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter v1.0.62

---

[JuliaFormatter v1.0.62] _{reported by reviewdog 🐶}

AbstractPPL.jl/test/varname/varname.jl

Line 151 in fbcde33

@test !is_dynamic(@varname(x[k = i]))

---

[JuliaFormatter v1.0.62] _{reported by reviewdog 🐶}

AbstractPPL.jl/test/varname/varname.jl

Lines 190 to 191 in fbcde33

	@test @varname(x[k = idx]) == @varname(x[k = 3])
	@test @varname(x[k = 2 * idx]) == @varname(x[k = 6])

[codecov]

Codecov Report

❌ Patch coverage is 55.62500% with 71 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.08%. Comparing base (76dc6df) to head (30bd5dc).

Files with missing lines	Patch %	Lines
ext/AbstractPPLDifferentiationInterfaceExt.jl	0.00%	36 Missing ⚠️
ext/AbstractPPLLogDensityProblemsExt.jl	0.00%	15 Missing ⚠️
src/ADProblems.jl	65.78%	13 Missing ⚠️
src/utils.jl	89.85%	7 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #155      +/-   ##
==========================================
- Coverage   84.51%   78.08%   -6.44%     
==========================================
  Files          10       14       +4     
  Lines         562      721     +159     
==========================================
+ Hits          475      563      +88     
- Misses         87      158      +71     

☔ 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.

[github-actions]

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

[gdalle]

@wsmoses you may want to take a look at this Enzyme ext, especially since the test doesn't fully cover it yet?

[wsmoses]

don't make a closure/capture the evaluator here, just call evaluator directly in autodiff with all of the arguments. Also don't add this extra shim to flatten/unflatten

[gdalle]

This 1000-LOC PR introduces a lot of duplication across backends, for reasons that I still don't fully understand. Again, if DI is insufficient in any way for Turing's purposes, feel free to point it out here and I'll work on it: JuliaDiff/DifferentiationInterface.jl#992

[gdalle]

Why flatten and unflatten when you could keep the structured representation?

[gdalle]

Why a Vector instead of whatever array type is appropriate (e.g. on GPU)?

[gdalle]

Some DI backends can go up to order 2

[gdalle]

Suggested change

	grad_buf = zeros(eltype(x0), length(x0))
	grad_buf = zero(x0)

[gdalle]

Suggested change

	fill!(p.gradient_buffer, 0.0)
	fill!(p.gradient_buffer, zero(eltype(x)))

[gdalle]

There are probably packages for flattening stuff more cleanly

[gdalle]

Why use a Vector to flatten e.g. a named tuple of GPU arrays?

[gdalle]

The unflattening doesn't have the same type as the initial input

[gdalle]

That is obviously not true, you (or Claude?) literally added that extension above

[gdalle]

The test code is duplicated across all backends, which is a bit of a shame

[gdalle]

Forgot to pick the correct review option

[penelopeysm]

As requested, I've not looked at the AD extensions. The biggest comments I have are

• Consider moving to another package please
• Docs would be really useful.

[penelopeysm]

Is there a reason why AbstractFloat rather than Real? A looser type bound of Real would allow for e.g. higher-order or sparse AD without affecting the common use case of Float64

[penelopeysm]

I think this argument shouldn't be a Symbol (notice that if you pass mode=:NotARealMode, it will do the Jacobian which is probably not what was intended). Either a custom type would be better

abstract type Mode end
struct Gradient <: Mode end
struct Jacobian <: Mode end

struct Foo{F<:Mode}
   ...
end

[penelopeysm]

Is it mandatory to re-export at the top level? Would it be sufficient to export from the submodule and call it a day?

[penelopeysm]

is there a reason to use lts instead of min which the original CI job uses?

[penelopeysm]

Is this dep actually used in src/? I couldn't see any usage. In which case, this doesn't seem like it's needed here or in Project.toml.

[penelopeysm]

Like if this was adtype::AbstractADType then I could see us needing the ADTypes dep. But right now that's not the case

[penelopeysm]

These methods can be handled with error hints too, see e.g. https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/5a3ea3845c90aa69d4d23720f5c299af25268a9a/DifferentiationInterface/src/init.jl#L1-L28 for some ideas perhaps?

[yebai]

Thanks for the review. Changes made:

AbstractFloat → Real — applied across src/ADProblems.jl and all AD extension files.

Removed explicit derivative mode — prepare no longer takes mode=:gradient/:jacobian. The derivative kind is
inferred from the prepared function’s prototype output:

• scalar output → use value_and_gradient
• vector output → use value_and_jacobian
• calling the wrong derivative API gives a direct ArgumentError.

Dropped AbstractPPL.ADProblems.dimension — the LDP-facing case is handled by LogDensityProblems.dimension.

Removed AbstractPrepared{:gradient/:jacobian} — prepared evaluators now subtype plain AbstractPrepared. The one remaining capability query uses an internal _supports_gradient trait instead.

LDP capabilities — avoided evaluating user code inside LogDensityProblems.capabilities(::AbstractPrepared).
Type-level capabilities still default conservatively to LogDensityOrder{0}() per the LDP convention; value-level methods expose gradient capability where the prepared object can actually support it.

Error handling — removed throwing fallback methods for prepare(adtype, ...), value_and_gradient, and
value_and_jacobian where they could cause method-ambiguity issues; missing backend guidance is handled with
register_error_hint in __init__. test_autograd’s missing-backend throw changed from ArgumentError to error().
The integer-vector throw(MethodError(...)) overloads on VectorEvaluator are intentionally kept because otherwise
AbstractVector{Int} silently matches the AbstractVector call method.

Checked → Validate — renamed throughout both evaluator types and their docstrings, including a note on the Trivial type-stability implication.

NT inputs in LDP ext — removed NamedTupleEvaluator-specific logdensity, dimension, and capabilities overloads; NT-backed evaluators should be used via value_and_gradient(prepared, values::NamedTuple) directly.

ADTypes dep — removed using ADTypes: ADTypes from src/; kept in [deps] since the extension files import from it.
Aqua’s stale-deps check is configured to ignore it.

Docs — added a dedicated evaluator/AD docs page covering structural vs AD preparation, vector inputs, NamedTuple
inputs, Jacobians, non-AD structural preparation, test_autograd, and supported backends.

Docstring typo (line 28), CI (lts → min, cleaner test command) — fixed.

---

Not changed:

• LDP as weak dep — kept as an extension because it is only needed for interop and does not need to be loaded for the core API.
• test_autograd generalisation — deferred; FiniteDifferences-only for now.
• CI continue-on-error — not redundant with fail-fast: false: the former prevents Enzyme failures from failing the overall workflow check; the latter prevents other matrix jobs from being cancelled. Both are needed.

[penelopeysm]

Could you (or Claude) explain why you are rejecting the parts of the review that you rejected?

[penelopeysm]

(I think it's important for future readers to understand why the decisions were made that way)

[yebai]

Thanks, Penny. Docs have been added. All other raised issues should be addressed or explained in the comments above.

As this will likely only be implemented within TuringLang, I don't think registering a new package for it would be worth
it. I created AbstractPPL specifically for functionalities like those in this PR.

EDIT: fixed a few edge cases for AbstractPPLDIExt (1) AutoReverseDiff needs to respect the provided compiled tag (2) AutoEnzyme need DI.Const for evaluator objects.

[penelopeysm]

Claude is also wrong about keeping this as a dep. If you need something that is only loaded in the extensions, it should go under weakdeps.

[yebai]

ADTypes is actually needed as a hard dependency for robust error handling -- it is on the fence, but I think it is useful as a hard dependency as it simplifies weak deps triggering, eg, ForwardDiff extension only needs using AbstractPPL, ForwardDiff.

[penelopeysm]

This function is not type-stable for NamedTuples which will probably translate into perf issues:

julia> using AbstractPPL

julia> @inferred AbstractPPL.Utils.unflatten_to!!((a=1, b=2, c=3), zeros(3))
ERROR: return type @NamedTuple{a::Float64, b::Float64, c::Float64} does not match inferred return type NamedTuple

julia> @inferred AbstractPPL.Utils.unflatten_to!!((a=1, b=[2, 3], c=3), zeros(4))
ERROR: return type @NamedTuple{a::Float64, b::Vector{Float64}, c::Float64} does not match inferred return type NamedTuple

It's very likely that you need a generated function to ensure type stability for NamedTuples.

Also tests to check type stability for the various inputs would probably be a good thing. There are loads of examples in DynamicPPL's VNT tests which might be helpful inspiration

[yebai]

Good catch, here, too.