Skip to content

Use AbstractPPL evaluator interface#255

Open
yebai wants to merge 9 commits into
mainfrom
abstractppl-dynamicppl-migration
Open

Use AbstractPPL evaluator interface#255
yebai wants to merge 9 commits into
mainfrom
abstractppl-dynamicppl-migration

Conversation

@yebai
Copy link
Copy Markdown
Member

@yebai yebai commented Apr 24, 2026

Replacing #249

Enzyme failures are not unrelated -- they fail on main

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>
@yebai yebai force-pushed the abstractppl-dynamicppl-migration branch from 639e031 to f689bd8 Compare April 24, 2026 15:36
yebai and others added 5 commits May 19, 2026 18:02
- 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>
`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>
@yebai yebai marked this pull request as ready for review May 20, 2026 13:12
yebai and others added 2 commits May 20, 2026 14:12
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant