Add Evaluators: prepare interface, and vectorisation utilities by yebai · Pull Request #157 · TuringLang/AbstractPPL.jl

yebai · 2026-04-30T12:44:45Z

Introduces AbstractPPL.Evaluators, a new submodule that provides the structural plumbing for preparing and calling model evaluators, with or without AD.

Core functions

prepare(problem, x; check_dims=true) / prepare(adtype, problem, x; check_dims=true) — structural and AD-aware preparation; AD-backend extensions implement the three-argument form
value_and_gradient!!(prepared, x) / value_and_jacobian!!(prepared, x) — stub interface for derivative computation; returned arrays may alias internal cache buffers
evaluate!!(evaluator, x) — extends the existing AbstractPPL.evaluate!! for Prepared, VectorEvaluator, and NamedTupleEvaluator

Evaluators

VectorEvaluator{CheckInput} — wraps a callable with a flat-vector input contract; CheckInput=false skips the per-call length check for AD hot paths
NamedTupleEvaluator{CheckInput} — same for NamedTuple inputs with a fixed prototype

Utilities

flatten_to!!(buf, x) / unflatten_to!!(x, buf) — round-trip vectorisation for Real, Complex, AbstractArray, Tuple, and NamedTuple; @generated NamedTuple unflatten preserves type
inferability

LogDensityProblems extension

logdensity, dimension, capabilities for Prepared and VectorEvaluator
logdensity_and_gradient copies the gradient out of the aliasable !! buffer

@generated

…and vectorisation utilities Introduces `AbstractPPL.Evaluators`, a new submodule providing structural plumbing for preparing and calling model evaluators, with or without AD. Core types and functions: - `Prepared{AD,E,C}` — wraps an evaluator with an AD backend type (for dispatch) and an optional backend-specific cache; exported from AbstractPPL - `prepare(problem, x; check_dims=true)` / `prepare(adtype, problem, x; check_dims=true)` — structural and AD-aware preparation; AD-backend extensions implement the three-argument form and thread `check_dims` to `VectorEvaluator{check_dims}` - `value_and_gradient!!(prepared, x)` / `value_and_jacobian!!(prepared, x)` — stub interface for derivative computation; returned arrays may alias internal cache buffers (callers copy if they need to retain past the next call) - `evaluate!!(evaluator, x)` — extends the existing AbstractPPL.evaluate!! for Prepared, VectorEvaluator, and NamedTupleEvaluator Evaluator shapes: - `VectorEvaluator{CheckInput}` — wraps a callable with a flat-vector input contract; CheckInput=false skips the per-call length check for AD hot paths - `NamedTupleEvaluator{CheckInput}` — same for NamedTuple inputs with a fixed prototype Vectorisation utilities: - `flatten_to!!(buf, x)` / `unflatten_to!!(x, buf)` — round-trip vectorisation for Real, Complex, AbstractArray, Tuple, and NamedTuple; @generated NamedTuple unflatten preserves type inferability LogDensityProblems extension (updated): - logdensity, dimension, capabilities for Prepared and VectorEvaluator - logdensity_and_gradient copies the gradient out of the aliasable !! buffer Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

yebai · 2026-04-30T12:47:25Z

@penelopeysm, this PR keeps only the vector-/NamedTuple-typed evaluator in #155 and removes all autograd-related extensions (which can be added later). This should be ready for a final look.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

codecov · 2026-04-30T12:50:12Z

Codecov Report

❌ Patch coverage is 86.71329% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.10%. Comparing base (76dc6df) to head (7ac3e3a).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
ext/AbstractPPLLogDensityProblemsExt.jl	0.00%	12 Missing ⚠️
src/evaluators/utils.jl	92.30%	6 Missing ⚠️
src/evaluators/Evaluators.jl	98.11%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #157      +/-   ##
==========================================
+ Coverage   84.51%   85.10%   +0.58%     
==========================================
  Files          10       13       +3     
  Lines         562      705     +143     
==========================================
+ Hits          475      600     +125     
- Misses         87      105      +18

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

- docs/src/evaluators.md: use `Prepared{<:AutoForwardDiff}` instead of `Prepared{AutoForwardDiff}` so dispatch matches the parametric instance type - test/evaluators/Evaluators.jl: rename DummyProblem → DummyModel, rename the testset and parameter names to drop ADProblems-era terminology Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

yebai · 2026-04-30T12:54:25Z

@shravanngoswamii can you also help review this?

github-actions · 2026-04-30T12:57:54Z

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

- Remove `export Prepared` and the `using .Evaluators: Prepared` re-import - Update docs example to import `Prepared` from `AbstractPPL.Evaluators` - Add a brief comment on the `evaluate!!` overloads Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…ties to order 0 - Dispatch all LDP methods on `Prepared{<:Any,<:VectorEvaluator}` so `NamedTupleEvaluator`-backed `Prepared` does not match (it cannot satisfy LDP's vector-input contract). - Default `capabilities` returns `LogDensityOrder{0}`; AD-backend extensions opt into `LogDensityOrder{1}` by overloading on their cache type. Without the overload, `logdensity_and_gradient` would hit the `value_and_gradient!!` stub and fail. - Test: demonstrate the backend opt-in pattern (TestADType overloads `capabilities`), and verify NamedTupleEvaluator-backed `Prepared` has no LDP methods defined (`dimension` throws, `capabilities` falls through to LDP's `nothing` default). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

penelopeysm · 2026-04-30T15:57:22Z

I'll take a more thorough look later. For my own edification, what's the relationship between this and #155? Is this meant to be "Part 1" of that PR (in which case that could be closed?) or something else?

yebai · 2026-04-30T16:04:05Z

This PR should supersede #155. Later PRs will add autograd extensions.

Reconstruct each leaf using `x`'s types rather than `eltype(buf)`, so a heterogeneous round-trip preserves `typeof(x)` instead of widening to the flat buffer's eltype. Add an opt-in `check_eltype` keyword that warns when `eltype(buf)` differs from `flat_eltype(x)`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The previous wording attributed the input-shape guarantee to AD libraries, but ForwardDiff/Mooncake do not validate user input. The actual contract is that the outer entry point validates `length(x)` once, and dual arrays derived from that `x` then flow through the AD loop without re-checking. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

- check_dims=false is an opt-in trust mode where the caller takes responsibility for length(x); the AD inner-loop motivation is now framed as a typical use, not a guaranteed validation chain. - The no-AD `prepare(problem, x)` form exists so callers can write `prepare(...)` uniformly without knowing whether a backend is loaded, and downstream primal-only consumers can accept the result either way. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

These helpers are unused inside this PR but are intended for AD-backend extensions to share, so each `value_and_gradient!!` / `value_and_jacobian!!` produces a uniform error rather than rolling its own. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

- test/run_extras.jl: drop the label→subdir mapping, use the label as the path directly, rename the included test file to main.jl - src/evaluators/utils.jl: - simplify flat_length tuple/namedtuple cases to sum(flat_length, ...) - replace getfield with x[name] in the _unflatten generator - reject Symmetric/Hermitian/Diagonal/Triangular/Tri/Bi-diagonal arrays up front rather than emitting broken round-trip results - src/evaluators/Evaluators.jl: fold integer-input rejection into the VectorEvaluator method body (parametrise on T, branch on T<:Integer); drop the unused first arg of _reject_integer_input - docs/src/evaluators.md: fix the f=...\\ncache=... indentation typo; expand the !! aliasing note with an example showing safe usage Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

yebai · 2026-05-01T10:20:53Z

Thanks for the careful review! All comments have been addressed.

test/run_extras.jl (6c9e4f2)

Label simplified to ext/logdensityproblems; the lookup table is gone. Renamed the included file to main.jl so the include is now include(joinpath(@__DIR__, label, "main.jl")) as suggested. CI
updated.

src/evaluators/utils.jl

flat_length tuple/namedtuple: now sum(flat_length, ...; init=0) (6c9e4f2).
_unflatten NamedTuple generator: getfield(x, ...) → x[...] (6c9e4f2).
Heterogeneous round-trip type loss (typeof(x2) != typeof(x)): unflatten_to!! now reconstructs leaves using x's types via convert(typeof(x), buf[i]) and similar(x). Round-trip preserves
typeof(x). Added an opt-in check_eltype::Bool=false keyword that warns when eltype(buf) !== flat_eltype(x). Regression tests added (d0ebb6c).
Structured arrays (Symmetric/Hermitian/Diagonal/Tri/Bi/SymTridiagonal/AbstractTriangular): rejected up front via a _StructuredArray union with a clear ArgumentError, rather than emitting broken
results. Cholesky/LU/QR were already rejected by the catch-all (not <:AbstractArray). Added a TODO noting that proper support for structured arrays and Cholesky factors (relevant to PPL covariance
parameters) is a follow-up (6c9e4f2, 88f1661).

src/evaluators/Evaluators.jl

Integer rejection: folded into the VectorEvaluator{true/false} method bodies, parametrising on T and branching on T <: Integer so the compiler optimises it away for non-integer inputs. The helper
now takes only x (no unused first arg) (6c9e4f2).
The unused _assert_supported_output / _assert_jacobian_output helpers: kept and given an explanatory comment — they're for AD-backend extensions to share so each value_and_gradient!! /
value_and_jacobian!! produces a uniform error rather than rolling its own (3ad1396).

docs/src/evaluators.md

Indentation typo in the prepare example fixed (6c9e4f2).
!! aliasing: expanded with a concrete example showing the copy(grad) pattern and the failure mode it prevents (6c9e4f2).
check_dims=false rationale rewritten — no longer attributes the input-shape guarantee to AD libraries; framed as an opt-in trust mode where the caller takes responsibility, with the AD inner-loop
scenario as a typical use (73c97e9, c2b84b2).
No-AD prepare(problem, x) rationale: now explains the use case — callers can write prepare(...) uniformly without knowing whether a backend is loaded, and downstream primal-only consumers (e.g.
log-density evaluation without gradients) accept the result either way (c2b84b2).

penelopeysm

just a couple more things

- Drop `init=0` from `flat_length` tuple/namedtuple sums and add explicit empty-case methods to mirror `flat_eltype`'s pattern. - Reject non-one-based arrays via `Base.require_one_based_indexing` (previously `copyto!(buf, offset, x, 1, n)` BoundsErrored on OffsetArrays silently). Checked at user-facing buf entry points and at the array dispatches in `_flatten_to!` / `_unflatten`. - Regression tests for both rejection paths. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

yebai · 2026-05-01T13:07:49Z

@penelopeysm / @shravanngoswamii should be ready to merge if you are happy.

shravanngoswamii · 2026-05-01T13:10:48Z

@penelopeysm / @shravanngoswamii should be ready to merge if you are happy.

If it's fine, can you please wait for the end of the day! I am happy to approve it afterwards, let me look more carefully at it.

penelopeysm · 2026-05-01T13:12:21Z

I can't, in good conscience, approve this PR, because I still don't agree with the direction. If you want to merge it, you have the permissions on GitHub to do so.

yebai · 2026-05-01T13:14:48Z

I can't, in good conscience, approve this PR, because I still don't agree with the direction. If you want to merge it, you have the permissions on GitHub to do so.

@penelopeysm, your input on Julia software engineering is appreciated.

The previous shape check relied on `typeof` equality, which treats `(b=zeros(2),)` and `(b=[1.0],)` as the same shape because both are `Vector{Float64}`. Add a structural `_shapes_match` recursion that compares `size` for array leaves (recursing into non-numeric eltypes) and rejects unsupported leaf types up front, matching the `Real`/`Complex`/`AbstractArray`/`Tuple`/`NamedTuple` contract used by the flatten/unflatten utilities. Document the supported leaves in the `NamedTupleEvaluator` docstring and the user-facing docs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Replace the `_StructuredArray` reject-list with a `FlattableArray = Union{Array,SubArray}` opt-in. This rejects `Adjoint`/`Transpose` (whose `similar` strips the wrapper and silently breaks the type round-trip) and any other custom `AbstractArray` whose round-trip we cannot guarantee. Also clarify the `check_eltype` warning to mention the `InexactError` that follows when narrowing conversions fail. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

yebai · 2026-05-01T23:18:46Z

All perspectives are appreciated. This PR introduces a self-contained evaluator interface that is reusable across DynamicPPL, Bijectors, and AdvancedVI, and cleanly separates primal evaluation from automatic differentiation interfaces, which will be introduced in subsequent PRs.

- src/evaluators/utils.jl: drop SubArray from FlattableArray allowlist; views silently lost their wrapper through similar(::SubArray) so the typeof round-trip invariant did not hold. - src/evaluators/Evaluators.jl: extend NamedTupleEvaluator docstring to mirror VectorEvaluator's CheckInput=false guidance, calling out that the prototype typeof check rejects dual/shadow leaves; narrow the AD-extension error hint to fire only when no extension has registered an AD-aware prepare method; add a brief comment on VectorEvaluator{false} noting the `T <: Integer` branch resolves at compile time. - ext/AbstractPPLLogDensityProblemsExt.jl: drop redundant instance-method capabilities delegations now that we rely on LogDensityProblems' built-in capabilities(x) = capabilities(typeof(x)). - test/evaluators/utils.jl: extend rejection test with a SubArray case; drop obsolete view round-trip edge case. - test/evaluators/Evaluators.jl: replace the internal _assert_namedtuple_shape direct-call asserts with public-call coverage. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

- src/evaluators/Evaluators.jl: collapse three trivial evaluate!! overloads into one Union method; remove unused _assert_jacobian_output and _assert_supported_output helpers (will be reintroduced by AD-extension PRs that actually call them); consolidate the integer-rejection compile-time rationale onto the helper; document the m.nargs >= 4 gate in the error hint. - src/evaluators/utils.jl: drop the single-use FlattableArray{T} = Array{T} alias and inline Array{...} at its 4 use sites; tighten the allow-list rationale onto the typeof round-trip contract; fold the heterogeneous round-trip note into the unflatten_to!! docstring and drop the verbose worked-example comment block (covered by the dedicated test). - ext/AbstractPPLLogDensityProblemsExt.jl: add a one-line WHY for the copy(grad) in logdensity_and_gradient (LDP requires a non-aliased gradient). - test/evaluators/utils.jl: trim the seven-element non-Array rejection loop to three representatives (structured wrapper, view, non-one-based) and drop the unused Diagonal/UpperTriangular imports. - test/run_extras.jl: collapse the two-step LABEL validation into a single membership check. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

github-actions Bot assigned yebai Apr 30, 2026

Format with JuliaFormatter (blue style)

0a3a935

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

format

d632db5

shravanngoswamii self-requested a review April 30, 2026 12:56

yebai and others added 2 commits April 30, 2026 14:02

penelopeysm reviewed May 1, 2026

View reviewed changes

shravanngoswamii reviewed May 1, 2026

View reviewed changes

Comment thread src/evaluators/Evaluators.jl Outdated

yebai and others added 6 commits May 1, 2026 11:03

Note TODO for proper structured-array / Cholesky support

88f1661

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

penelopeysm reviewed May 1, 2026

View reviewed changes

Comment thread src/evaluators/utils.jl Outdated

Comment thread src/evaluators/utils.jl

shravanngoswamii requested changes May 1, 2026

View reviewed changes

Comment thread src/evaluators/utils.jl

Comment thread src/evaluators/utils.jl Outdated

yebai requested a review from shravanngoswamii May 1, 2026 23:19

shravanngoswamii reviewed May 2, 2026

View reviewed changes

Comment thread src/evaluators/Evaluators.jl

Comment thread src/evaluators/Evaluators.jl

Comment thread src/evaluators/Evaluators.jl

Comment thread src/evaluators/utils.jl Outdated

shravanngoswamii approved these changes May 4, 2026

View reviewed changes

yebai merged commit ba8d21c into main May 4, 2026
13 checks passed

yebai deleted the evaluators branch May 4, 2026 17:17

yebai mentioned this pull request May 4, 2026

Add evaluator interface with AD backend extensions #155

Closed

3 tasks

Conversation

yebai commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Core functions

Evaluators

Utilities

LogDensityProblems extension

Uh oh!

yebai commented Apr 30, 2026

Uh oh!

codecov Bot commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

yebai commented Apr 30, 2026

Uh oh!

github-actions Bot commented Apr 30, 2026

Uh oh!

penelopeysm commented Apr 30, 2026

Uh oh!

yebai commented Apr 30, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yebai commented May 1, 2026

Uh oh!

penelopeysm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yebai commented May 1, 2026

Uh oh!

shravanngoswamii commented May 1, 2026

Uh oh!

penelopeysm commented May 1, 2026

Uh oh!

yebai commented May 1, 2026

Uh oh!

Uh oh!

Uh oh!

yebai commented May 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yebai commented Apr 30, 2026 •

edited

Loading

codecov Bot commented Apr 30, 2026 •

edited

Loading