Use AbstractPPL evaluator interface#255
Open
yebai wants to merge 9 commits into
Open
Conversation
Replace `DifferentiationInterface` with `AbstractPPL.prepare` / `AbstractPPL.value_and_gradient` throughout. Key changes: - `_prepare_gradient` / `_value_and_gradient!` now wrap an `AbstractPPL` prepared evaluator in a `_VIGradPrep` struct that holds an `aux_ref` so auxiliary inputs can be swapped without re-preparing. - `DynamicPPLModelLogDensityFunction` stores `model_ref` and `loglikeadj_ref` as `Ref`s; `subsample` mutates them in-place instead of creating a new struct via `@set`, keeping the prepared evaluator valid across subsampling steps. - Drop second-order (Hessian) support; `use_hessian=true` now warns and is ignored. - Pin dev branches via `[sources]`: `AbstractPPL@evaluator-interface` and `DynamicPPL@adproblems-interface`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
639e031 to
f689bd8
Compare
- Switch `value_and_gradient` → `value_and_gradient!!` per AbstractPPL 0.15.
- In AdvancedVI core, give every `AbstractPPL.Prepared{<:AbstractADType,<:VectorEvaluator}`
a `LogDensityProblems` interface (order-1 fallback) so any AD-backed prep
acts as a `LogDensityProblem` without backend-specific wiring.
- In `AdvancedVIMooncakeExt`, promote Mooncake-prepared evaluators to order-2
and add `logdensity_gradient_and_hessian` via forward-over-reverse Mooncake.
- Simplify `DynamicPPLModelLogDensityFunction` to delegate LDP calls to its
inner prep; `capabilities` reads off the prep's own capability so the Hessian
branch is exposed exactly when the backend supports it.
- Bump compat: AbstractPPL 0.15, DynamicPPL 0.42 (with branch pin until
release); add Bijectors branch pin in `test/` for the same reason.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`git diff main --stat` was full of `return` keyword and `*` spacing adjustments unrelated to the AbstractPPL 0.15 switch. Restoring those files to `main`'s state shrinks the review surface to just the load-bearing changes (AbstractPPL/DynamicPPL bump, LDP+Hessian wiring, stale-DI comment fixes). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Keeps this branch focused on the AbstractPPL evaluator-interface migration: the LDP fallback methods on `Prepared`, the Mooncake forward-over-reverse Hessian, and the DynamicPPLModelLogDensityFunction LDP-delegation refactor have been moved to the `ldp-on-prepared-hessian` branch (branched from `main` and carrying both the migration and the feature work). What stays here: - Switch DI → AbstractPPL.prepare / value_and_gradient!! in core. - DynamicPPLModelLogDensityFunction uses AbstractPPL.prepare instead of DI.prepare_gradient (Hessian path is dropped as in the original migration commit; `use_hessian=true` warns and is ignored). - Compat bumps: AbstractPPL@0.15, DynamicPPL@0.42, plus Bijectors source pin in `test/`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Pins AbstractPPL to `hg/hessian-order`, which adds `prepare(adtype, f, x; order=2)` and `value_gradient_and_hessian!!`. `DynamicPPLModelLogDensityFunction` goes back to main's `prep_grad` + `prep_hess` shape (default `use_hessian=true`), so the diff reads as a clean DI → AbstractPPL swap rather than a redesign. `use_hessian` falls back to gradient-only when the AD backend refuses `order=2` (MethodError only; other errors propagate). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Wrapping the prepared evaluator and the aux `Ref` in a struct was just working around the lack of mutable state on `AbstractPPL.Prepared`. AbstractPPL 0.15's `prepare(...; context=tuple)` keeps the tuple on the evaluator and threads it through every call, so passing `Ref(aux)` in the context stores the ref inside the prep itself; callers mutate `prep.evaluator.context[1][]` before each evaluation. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1ea6829 to
79db2fb
Compare
`loglikeadj_ref` now tracks the input numeric type via a `LogLikeAdj<:Real`
struct parameter, restoring `Float32`/`BigFloat`/dual support that the
hard-pinned `Ref{Float64}` quietly dropped. `subsample` computes the
adjustment in the field's eltype to preserve precision across minibatches.
The block comment now also flags why `model_ref::Ref{Any}` cannot be
tightened (subsample-widened `defaults` NamedTuple), so a future reader
doesn't try.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
AbstractPPL 0.15.1 and Bijectors 0.16.0 are now released. The DynamicPPL `adproblems-interface` branch has been updated to require Bijectors 0.16, so keeping the DynamicPPL source pin no longer pulls in an older Bijectors. Restrict the package and test envs to Bijectors 0.16; relax docs to "0.15, 0.16" since NormalizingFlows 0.2.2 still requires 0.15.x. 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
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.
Replacing #249
Enzyme failures are not unrelated -- they fail on
main