Part 2 of autograd backend: Mooncake extention#160
Open
yebai wants to merge 18 commits into
Open
Conversation
Layer the DifferentiationInterface AD extension onto the simplified evaluator
interface introduced on the `evaluators` branch. This is the catch-all path
for any `<:AbstractADType` backend without a native AbstractPPL extension.
- ext/AbstractPPLDifferentiationInterfaceExt.jl: DI-backed `prepare`,
`value_and_gradient!!`, and `value_and_jacobian!!` for vector-input
evaluators. The evaluator is passed as a `DI.Constant` so DynamicPPL-style
problem state stays constant across calls. Compiled `AutoReverseDiff{true}`
takes the one-argument tape path (the `DI.Constant` route would invalidate
the compiled tape across calls). Empty inputs short-circuit before
`DI.prepare_*` since many backends fail on length-zero arrays.
- Project.toml: DifferentiationInterface added as a weakdep with extension
trigger and compat bound.
- test/autograd_tests.jl: shared problem definitions and `run_autograd_tests`
entry point with separate gradient / jacobian / empty-input helpers.
- test/ext/differentiationinterface/: isolated test env using a local
`DummyADType <: AbstractADType` to exercise the catch-all dispatch without
pulling in a real AD package.
- .github/workflows/CI.yml: extend the ext matrix with
`ext/differentiationinterface` alongside the existing
`ext/logdensityproblems`.
- test/run_extras.jl: register the new label.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The upstream `Simplify Evaluators module after review` commit dropped `_assert_supported_output` / `_assert_jacobian_output` from `Evaluators.jl`, expecting the AD-extension PRs that actually call them to bring them back. Move both helpers into the DI extension, the sole caller after rebase. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replace `_assert_supported_output` / `_assert_jacobian_output` with an `if y isa Number / elseif y isa AbstractVector / else throw` cascade. The jacobian assertion was unreachable after the scalar branch, and the supported-output helper had a single call site; inlining the remaining check as the `else` arm is clearer than two named helpers. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Expose shared AD backend fixtures through TestResources so extension test environments can reuse the same cases without including files from test/. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Exercise the DifferentiationInterface catch-all with a real backend instead of a stub so the extension test covers the integration path used by downstream AD packages. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
… extension Replace `prepare(adtype, problem, x)`'s two-step `prepare` + `VectorEvaluator` wrap with a single delegated call to `AbstractPPL.prepare(problem, x; check_dims)`, eliminating a latent double-wrap when `check_dims=false`. Move `TestResources` (callable problems, test cases, runner) out of `src/` and into a new `AbstractPPLTestExt` package extension, so `Test` is loaded only when downstream code opts in. Split the unified `TestCase` into `ValueCase` / `ErrorCase`, drop the unused `:namedtuple` fixtures, and store the error operation as a callable instead of a `:call`/`:gradient` symbol. Flatten the DI extension test file to two `run_testcases(Val(...); ...)` calls. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- value_and_gradient!! / value_and_jacobian!! now check arity before the
length-zero short-circuit, so empty-input vector-valued / scalar-valued
functions surface the correct ArgumentError instead of silently returning
the wrong shape. Uses Val(0) as the cache sentinel for empty inputs.
- New AbstractPPLDifferentiationInterfaceLogDensityProblemsExt advertises
LogDensityOrder{1} for any DI-prepared evaluator, registering the method
in __init__ to side-step extension precompile-visibility constraints.
- Empty-input arity errors and the new capability advertisement are
exercised in the DI test environment, which now also loads LDP.
- Project.toml [compat] alphabetised.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Drop the triple extension and have AbstractPPLLogDensityProblemsExt
advertise LogDensityOrder{1} for any Prepared{<:Any,<:VectorEvaluator}.
The contract is now uniform: anything from `prepare(adtype, …)` claims
gradient capability; a bare VectorEvaluator from `prepare(problem, x)`
stays at order 0. Backends that don't implement value_and_gradient!! or
that wrap a vector-output function surface a runtime error at call time
rather than via a capability downgrade.
Documents the rule in docs/src/evaluators.md.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…tras Aqua's persistent_tasks check spawns a wrapper subprocess that runs Pkg.precompile() on a package depending on AbstractPPL. On Julia 1.10 this hits "Declaring __precompile__(false) is not allowed in files that are being precompiled" inside the wrapper's extension precompile path. The dedicated Ext CI jobs load and exercise every extension on min Julia and pass, so this is a Julia 1.10 / Aqua interaction, not a defect in our extensions. Re-enable when min is bumped past 1.10. The CI extras step also had a redundant --project=. that was overridden by run_extras.jl's own Pkg.activate; drop it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Mark generate_testcases / run_testcases as public (1.11+) so downstream AD-backend packages can reuse the conformance suite without reaching into private API. - Note in the Prepared docstring that the two-arg constructor is for backends that allocate fresh storage per call. - Expand the _ad_output_arity error message to call out matrix/tuple outputs explicitly and recommend flattening. - Add a MethodError hint on value_and_gradient!! / value_and_jacobian!! that fires when no AD backend extension is loaded — also surfaces through LogDensityProblems.logdensity_and_gradient, which delegates to value_and_gradient!!. Suppressed once any backend registers a method, where the standard candidate list is more informative. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
|
AbstractPPL.jl documentation for PR #160 is available at: |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #160 +/- ##
===========================================
- Coverage 85.10% 69.42% -15.69%
===========================================
Files 13 15 +2
Lines 705 870 +165
===========================================
+ Hits 600 604 +4
- Misses 105 266 +161 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
- Drop AbstractPPLLogDensityProblemsExt and its test launcher; downstream
packages can define LDP methods on `Prepared` directly when needed.
- Add `_check_vector_length` to enforce dim mismatch on
`value_and_gradient!!`/`value_and_jacobian!!` before the `Val(0)`
short-circuit and DI call (matches `VectorEvaluator{true}` behaviour).
- Unify gradient/jacobian DI prep into `_prepare_di(prep, ...)`; the
`AutoReverseDiff{true}` carve-out now also fixes Jacobian tape recording.
- Tidy DI tests: remove LDP block, drop redundant constants and edge cases,
tighten gradient-of-vector-valued case to `ArgumentError`.
- Document prep-time `evaluator(x)` call and reserved `Val{:vector}` /
`Val{:edge}` testcase keys.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The bare `VectorEvaluator` rejected integer vectors with a clear `ArgumentError`, but the DI derivative entry points went straight to DI prep and surfaced an opaque `PreparationMismatchError`. Apply the same compile-time `T <: Integer` check before dispatching to DI, and add matching `:edge` testcases. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mooncake AD-backend extension built on the evaluator interface, with the
shared conformance suite extended to cover NamedTuple inputs and empty-input
arity errors. Squashed from prior incremental commits:
- AbstractPPLMooncakeExt: cache reuse, scalar/vector dispatch, NamedTuple
inputs via VectorEvaluator/NamedTupleEvaluator wrappers; integration test
in test/ext/mooncake.
- Evaluators._ad_output_arity: lift the duplicated `Union{Number,
AbstractVector}` output check from both extensions into one helper that
returns `:scalar` / `:vector` for downstream dispatch.
- Empty-input arity tagging (`Val(:scalar)` / `Val(:vector)`) so the
empty-input fast path raises the same "requires a scalar/vector-valued
function" error as the DI path instead of silently succeeding.
- AbstractPPLTestExt: add `Val(:namedtuple)` group (one ValueCase + one
ErrorCase); tighten regex assertions on the existing arity-mismatch cases.
- check_dims threaded through the inner `prepare` call so AD hot paths can
skip per-call shape checks.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`_check_ad_input(evaluator, x)` in `Evaluators` replaces the duplicated
`T <: Integer` rejection plus length check that appeared at six AD entry
points (two in the DI extension, four in Mooncake). Compile-time `T`
elision is preserved.
Move `generate_testcases(::Val{:namedtuple})` and
`run_testcases(::Val{:namedtuple})` to sit alongside the `:vector` and
`:edge` definitions so the file reads generate-then-run for all three
groups.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- `:cache_reuse` conformance group in `AbstractPPLTestExt` drives
`value_and_{gradient,jacobian}!!` three times per case against a single
`prepared` evaluator to catch backend cache corruption between calls.
- `_assert_namedtuple_shape` always validates regardless of `CheckInput`;
the `{false}` callable opts out by not calling here. Fixes a silent
wrong-shape gradient when `prepare(...; check_dims=false)` was used with
a NamedTuple input whose nested array sizes differ from the prototype.
Mooncake regression test added.
- DI ext sub-environment now also loads `ReverseDiff` and exercises
`AutoReverseDiff(compile=true)` against the conformance suite, covering
the `_prepare_di(::AutoReverseDiff{true}, …)` compiled-tape path.
- Trimmed verbose `check_dims` clarifications in docstrings and
`docs/src/evaluators.md` to one sentence each.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Builds on the evaluator interface from #158 to add a Mooncake AD-backend extension and a shared, reusable AD conformance test suite.
AbstractPPLMooncakeExt— supportsAutoMooncake(reverse) andAutoMooncakeForwardfor bothVectorEvaluatorandNamedTupleEvaluatorinputs. Cache is arity-tagged (MooncakeCache{:scalar|:vector,…}) so wrong-arity calls raise the same"requires a scalar/vector-valued function"errors as the DI extension instead of failing inside Mooncake.AbstractPPLTestExt—generate_testcases/run_testcasesfor groups:vector,:edge,:namedtuple,:cache_reuse. Both DI and Mooncake extension tests reuse this suite; downstream backend packages can do the same.Evaluators._ad_output_arityandEvaluators._check_ad_inputlift duplicated arity / input-validation code from both extensions.value_and_{gradient,jacobian}!!now validate input length and reject integer vectors before invoking DI, matching the bareVectorEvaluatorcontract.Breaking changes (version bump 0.14.3 → 0.15)
AbstractPPLLogDensityProblemsExtis removed. Downstream packages that calledLogDensityProblems.logdensity(p::Prepared, x)should callp(x)directly;LogDensityProblems.logdensity_and_gradient(p, x)users should callvalue_and_gradient!!(p, x). The LDP integration was thin and downstream packages can wrap aPreparedin their own LDP-shaped type if they want the legacy API.Correctness fix in this PR
prepare(...; check_dims=false)with aNamedTupleinput previously made the publicvalue_and_gradient!!skip its shape check (because_assert_namedtuple_shape(::NamedTupleEvaluator{false}, _)was a no-op). Mooncake caches are tied to nested array sizes, so a mismatched leaf would silently return a wrong-shape gradient._assert_namedtuple_shapenow always validates regardless ofCheckInput; the{false}evaluator opts out only on the inner-evaluator hot path. New regression test intest/ext/mooncake/main.jl.Test plan
AutoForwardDiffandAutoReverseDiff(compile=true)exercise:vector,:cache_reuse,:edge.AutoMooncakeandAutoMooncakeForwardexercise:vector,:namedtuple,:cache_reuse,:edge, pluscheck_dims=falsecache-safety regression.🤖 Generated with Claude Code