From f855d4e3dbe578bde461e906c7dd89e3c14e5662 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Tue, 28 Jul 2026 22:13:35 -0400 Subject: [PATCH 1/2] Take HVPs through DI's second-order operator --- CHANGELOG.md | 67 +++++--- docs/src/10-getting-started.md | 18 ++- docs/src/15-gpu.md | 8 +- ext/DynamicPPLExt.jl | 34 +++- ext/LogDensityProblemsExt.jl | 36 ++++- src/DEER/DEER.jl | 64 +++++++- src/interface.jl | 274 +++++++++++++++++++------------- test/test-ADType-Slots.jl | 231 +++++++++++++++++++++------ test/test-Turing-Integration.jl | 86 ++++++++++ 9 files changed, 628 insertions(+), 190 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6dbf3aa..8b29749 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,24 +23,39 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 its own `hvp` / `hvp_batch` does not need it (#52). With no sampler `backend`, a batched HVP is derived from the model's own `hvp` backend. - A `logdensity_batch` given without a `grad_logdensity_batch` now has the - batched gradient derived for it, from the gradient slot's backend if it has - one and the sampler's otherwise, rather than leaving the batched DEER path - switched off (#52). `hvp_batch` can be a backend in that case too, and - differentiates the derived gradient; with no backend anywhere to derive from, - it raises when sampling starts. + batched gradient derived for it when `grad_logdensity` is a backend, rather + than leaving the batched DEER path switched off (#52). `hvp_batch` can be a + backend in that case too, and differentiates the derived gradient. +- An HVP backend over an AD-derived gradient is now taken as true second-order + AD, `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed + to `DI.hvp`, instead of an outer AD pass over the prepared DI gradient (#37). + That nesting dropped out of its preparation as soon as the outer pass pushed + tangents in, so the composed operator is both what was asked for and cheaper: + around 10x fewer allocations for a `logdensity`-only model. A backend over a + hand-written gradient still differentiates that gradient once, as before. +- `hvp` / `hvp_batch` accept a `SecondOrder` with both halves honoured, meaning + the log-density is differentiated twice and the gradient slot is not the inner + pass. Previously the inner half was silently discarded and only the outer used. + This is the one AD route to an HVP for a Turing or LogDensityProblems model, + whose gradient arrives already prepared and so cannot be differentiated again. +- The `DensityModel` constructors in `DynamicPPLExt` and `LogDensityProblemsExt` + forward `logdensity_batch`, `grad_logdensity_batch` and `hvp_batch`, so a + Turing or LogDensityProblems model can reach the batched DEER path. Neither + provides a batched log-density, so `logdensity_batch` has to be written by hand. - Adds `JuliaFormatter` testing which was forgotten (#60). ### Fixed - The reverse-on-grad HVP path differentiated with `DI.inner(backend)` while - its strategy was routed on `DI.outer(backend)`, so an `hvp` or `backend` - given as a `DifferentiationInterface.SecondOrder` ran the wrong half of the - pair. Both paths now take the outer, which is the pass being run--the - gradient slot is the inner one. Unwrapping to the outer half happens before - the backend-specific normalization hooks are dispatched on, so a - `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` and gets its mode - and function annotation pinned rather than running as a bare `AutoEnzyme()` - (which aborts on GPU). + its strategy was routed on `DI.outer(backend)`, so a + `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A + `SecondOrder` now goes to the true second-order path instead of either + strategy, and the half-selecting helpers it still uses agree: normalization + applies to the outer, the pass whose mode the backend extensions care about. + Unwrapping to that half happens before the normalization hooks are dispatched + on, so a `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` and gets + its mode and function annotation pinned rather than running as a bare + `AutoEnzyme()` (which aborts on GPU). ### Changed @@ -50,11 +65,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 reverse-on-grad path (#38). - Because a `logdensity_batch` without a `grad_logdensity_batch` now has the batched gradient derived rather than switching the batched DEER path off, a - model in that shape runs the batched update where it used to run the unbatched - one, and AD is applied to its `logdensity_batch`. On GPU that subjects a - function nothing was differentiating before to the backend's restrictions - (`pmcmc_*` wrappers for Enzyme). Supply `grad_logdensity_batch` to keep AD out - of the batched path. + model whose `grad_logdensity` is a backend runs the batched update where it + used to run the unbatched one, and AD is applied to its `logdensity_batch`. On + GPU that subjects a function nothing was differentiating before to the + backend's restrictions (`pmcmc_*` wrappers for Enzyme). Supply + `grad_logdensity_batch` to keep AD out of the batched path. A model with a + hand-written `grad_logdensity` is unaffected: nothing derives a batched + gradient for it, so the batched path stays off as before. +- `ParallelMALASampler`'s `backend` no longer derives a batched gradient, only + Hessian-vector products. It could previously switch the batched DEER path on + for a model with a hand-written gradient, which made a keyword that reads as + an HVP fallback decide which update path ran and put AD on a + `logdensity_batch` the user had not opted into differentiating. Models that + relied on that should pass `grad_logdensity_batch` explicitly, or a backend in + `grad_logdensity` for one to be derived from. +- Both batched derivative slots now require `logdensity_batch`, which the + batched update evaluates directly, and the constructor rejects them without + one. A callable `grad_logdensity_batch` or `hvp_batch` supplied on its own + used to be accepted and then silently ignored. `logdensity_batch` alone is + still valid and still used to score whole trajectories. +- An `hvp_batch` that reaches sampling with no batched gradient to pair it with + now raises rather than silently falling back to the unbatched update. ### Removed diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index 6a9f4b4..516ad81 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -39,9 +39,23 @@ model = DensityModel(logp, AutoEnzyme(), 2; param_names=[:x1, :x2]) Backends are turned into prepared [DifferentiationInterface](https://github.com/JuliaDiff/DifferentiationInterface.jl) callables when sampling starts, and that preparation is reused for the rest of the run. Hand-written and AD-derived slots mix, so an analytical gradient with `hvp=AutoForwardDiff()` is fine. -A backend in `hvp` differentiates whatever the gradient slot holds; it is not a second derivative of `logdensity`. Over an AD-derived gradient that composition is second-order AD, and over a hand-written one it is a single AD pass across your own code. The same goes for the batched pair, and a `logdensity_batch` supplied without a `grad_logdensity_batch` has the batched gradient derived for it — one gradient of `sum(logdensity_batch(X))`, which is the stacked per-column gradients only because columns are independent, so `logdensity_batch` must not couple them. +### What a backend in `hvp` differentiates -`backend` on [`ParallelMALASampler`](@ref) is the fallback derivative source for whatever the model did not bring: Hessian-vector products, and the batched gradient in the case just described. A model that supplies its own can leave it out. Note that passing one can switch the batched DEER path on, which also puts AD on your `logdensity_batch` — see the [GPU page](15-gpu.md) for when that matters. +How the HVP is taken depends on what the gradient slot holds: + +| `grad_logdensity` | `hvp` | what runs | +|---|---|---| +| callable | backend | one AD pass over your gradient | +| backend | backend | `SecondOrder(hvp, grad_logdensity)` on `logdensity` | +| either | `SecondOrder(...)` | that pair on `logdensity`, gradient slot unused | + +So a hand-written gradient with `hvp=AutoForwardDiff()` differentiates *your* code once, whereas a log-density-only model gets genuine second-order AD through DI's own second-order operator rather than one AD call nested inside another. Passing a `SecondOrder` yourself always means the log-density is differentiated twice — both passes are named, so the gradient slot is not the inner one even when you wrote it by hand. That is also the only AD route to an HVP for a Turing or LogDensityProblems model, whose gradient comes pre-prepared and cannot be differentiated again. + +Which combinations actually work is up to the backends, and second-order support is the thinnest part of the ecosystem. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model; `AutoMooncake` does not, in either direction — it has no reverse-over-reverse, and its gradient will not accept an outer pass's tangents, so a `SecondOrder` with a forward outer fails too. Write `grad_logdensity` out by hand for Mooncake and the single pass over it works normally. None of the second-order combinations work on GPU yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37) and the [GPU page](15-gpu.md)). + +The batched pair works the same way, on `sum(logdensity_batch(X))`. Its gradient is the stacked per-column gradients only because the columns are independent, so `logdensity_batch` must not couple them. A `logdensity_batch` given without a `grad_logdensity_batch` has one derived when `grad_logdensity` is a backend; with a hand-written gradient the batched path stays off and the unbatched update covers it. Both batched derivative slots require `logdensity_batch`, but `logdensity_batch` on its own is fine — it is also used to score a whole trajectory at once. + +`backend` on [`ParallelMALASampler`](@ref) is the fallback source of Hessian-vector products, for a model that brings no `hvp` / `hvp_batch` of its own. That is all it does: it never supplies a gradient, so it cannot change which update path runs. A model carrying its own HVPs can leave it out. --- diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 9153303..2a16cf3 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -233,10 +233,10 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You only supply `gradlogp` / `grad_logdensity_batch`.** The sampler builds the HVP by differentiating your gradient — either a forward-mode pushforward of `gradlogp` ([`ForwardOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), the default for most backends) or a reverse-mode gradient of `x -> dot(gradlogp(x), v)` ([`ReverseOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), used for `AutoMooncake` and `AutoZygote`). This is the **AD-HVP fallback**, and it is what the logistic-regression example above uses. !!! warning "Log-density-only models on GPU" - `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU. The AD-HVP fallback would then differentiate an AD-derived gradient, i.e. second-order AD, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). Write `gradlogp` out by hand for DEER. The sequential samplers only need the gradient, so log-density-only models are fine there. + `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU. The HVP is then true second-order AD — `SecondOrder(hvp_backend, grad_backend)` on your log-density — which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it. The sequential samplers only need the gradient, so log-density-only models are fine there. -!!! note "A `logdensity_batch` without a `grad_logdensity_batch` puts AD on it" - The batched DEER path needs both halves, so a `logdensity_batch` supplied on its own has the batched gradient derived by AD — from the gradient slot's backend if it has one, otherwise the sampler's `backend`. That is convenient on CPU, but on GPU it means the backend's restrictions above now apply to your `logdensity_batch`, even if every derivative you wrote by hand avoided them. Supply `grad_logdensity_batch` (and `hvp_batch`) to keep AD out of the batched path, or leave `logdensity_batch` off and let DEER run the unbatched update. +!!! note "A hand-written gradient keeps AD out of the batched path" + The batched DEER path needs a batched gradient as well as a `logdensity_batch`. One is derived from `logdensity_batch` only when `grad_logdensity` is a backend — i.e. only for a model that already opted into AD over its own code. With a hand-written `grad_logdensity` and no `grad_logdensity_batch`, the batched path simply stays off and the unbatched update runs, so nothing you wrote by hand is ever handed to AD behind your back. The sampler's `backend` does not change this; it only ever supplies HVPs. ### When the fallback is the right call @@ -249,7 +249,7 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **The HVP has a clean closed form.** Quadratic priors, Gaussian likelihoods, GLMs (logistic, Poisson, probit) — the second derivative is a known function of intermediate quantities you already compute in `gradlogp`. A few extra lines and you skip the AD pipeline entirely. - **Performance matters and the AD compile is heavy.** Enzyme and Mooncake both pay a one-shot compilation cost on the user's gradient. For long-running chains this amortizes, but for many short runs the analytical HVP wins. -- **You're hitting AD-backend-specific GPU restrictions.** The [Enzyme limitations](#2-enzyme-on-gpu-currently-needs-pmcmc_matmul-pmcmc_dot-pmcmc_dotsum) above (`pmcmc_*` wrappers, staged broadcasts) only matter when the AD backend is invoked. Supplying analytical HVP sidesteps them — your `gradlogp` and `hvp` can use plain `*`, `dot`, `sum`, and the sampler's `backend` can be omitted entirely because no AD is ever invoked. To keep it that way, supply `grad_logdensity_batch` alongside `logdensity_batch`: see the note below. +- **You're hitting AD-backend-specific GPU restrictions.** The [Enzyme limitations](#2-enzyme-on-gpu-currently-needs-pmcmc_matmul-pmcmc_dot-pmcmc_dotsum) above (`pmcmc_*` wrappers, staged broadcasts) only matter when the AD backend is invoked. Supplying analytical HVP sidesteps them — your `gradlogp` and `hvp` can use plain `*`, `dot`, `sum`, and the sampler's `backend` can be omitted entirely because no AD is ever invoked. A model whose derivatives are all hand-written never reaches AD by any route, batched path included. - **You can reuse intermediates between gradient and HVP.** When `hvp` shares $X\beta$, $\sigma(X\beta)$, or similar with the gradient computation, an analytical version can be both faster *and* shorter than what AD produces. ### Same example with analytical HVP diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index 0e5e92a..cabc650 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -8,7 +8,8 @@ using FlexiChains: FlexiChain, VarName, VNChain, SymChain using LogDensityProblems: LogDensityProblems """ - DensityModel(turing_model::DynamicPPL.Model; ad_backend, hvp=nothing) + DensityModel(turing_model::DynamicPPL.Model; ad_backend, hvp=nothing, + logdensity_batch=nothing, grad_logdensity_batch=nothing, hvp_batch=nothing) Convenience constructor: wraps a DynamicPPL/Turing `@model` directly as a `DensityModel`, automatically extracting parameter names and wiring up gradient @@ -17,6 +18,20 @@ computation via DynamicPPL's `adtype` interface. Requires `DynamicPPL` and `LogDensityProblems` to be loaded (these are the weak-dependency triggers for this extension), plus any AD backend that is used. +`ad_backend` is DynamicPPL's own `adtype`, not a `DensityModel` slot: it goes to +the `LogDensityFunction` that fills the log-density and gradient slots, which is +why it takes a backend only and never a callable. Everything else here is a +`DensityModel` slot forwarded unchanged, so `hvp` and the batched slots take a +callable or a backend as they do on the main constructor. + +For `ParallelMALASampler`, which also needs an HVP, pass either a callable or a +`DifferentiationInterface.SecondOrder` — the latter differentiates the +log-density twice, bypassing DynamicPPL's gradient. A plain backend does not +work here: it would differentiate the gradient `ad_backend` produced, whose AD +preparation rejects the tangents an outer pass pushes through it. DynamicPPL +supplies no batched log-density either, so `logdensity_batch` has to be written +out by hand to reach the batched DEER path. + # Example ```julia using Turing, ParallelMCMC, FlexiChains, ForwardDiff @@ -31,7 +46,14 @@ chain = sample(model, AdaptiveMALASampler(0.3; n_warmup=500), 2_000; chain_type=FlexiChains.VNChain, discard_warmup=true, progress=true) ``` """ -function ParallelMCMC.DensityModel(turing_model::DynamicPPL.Model; ad_backend, hvp=nothing) +function ParallelMCMC.DensityModel( + turing_model::DynamicPPL.Model; + ad_backend, + hvp=nothing, + logdensity_batch=nothing, + grad_logdensity_batch=nothing, + hvp_batch=nothing, +) # Sample in linked/unconstrained space and let DynamicPPL provide the gradient. ld = DynamicPPL.LogDensityFunction( turing_model, @@ -40,7 +62,13 @@ function ParallelMCMC.DensityModel(turing_model::DynamicPPL.Model; ad_backend, h adtype=ad_backend, ) # Requires LogDensityProblemsExt to be loaded - return ParallelMCMC.DensityModel(ld; hvp=hvp) + return ParallelMCMC.DensityModel( + ld; + hvp=hvp, + logdensity_batch=logdensity_batch, + grad_logdensity_batch=grad_logdensity_batch, + hvp_batch=hvp_batch, + ) end ###################### diff --git a/ext/LogDensityProblemsExt.jl b/ext/LogDensityProblemsExt.jl index a1c3885..c940a5e 100644 --- a/ext/LogDensityProblemsExt.jl +++ b/ext/LogDensityProblemsExt.jl @@ -4,7 +4,8 @@ using ParallelMCMC using LogDensityProblems: LogDensityProblems """ - DensityModel(ld; param_names=nothing, hvp=nothing) + DensityModel(ld; param_names=nothing, hvp=nothing, + logdensity_batch=nothing, grad_logdensity_batch=nothing, hvp_batch=nothing) Construct a `DensityModel` from any object implementing the [LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface. @@ -19,7 +20,18 @@ The optional `param_names` keyword accepts a collection of parameter names that for the columns of the returned `FlexiChain` object. If omitted, a single vector-valued parameter named `:x` will be chosen, unless you also pass `param_names` to `sample(...)`. -The `hvp` keyword argument is forwarded to the main `DensityModel` constructor. +`hvp` and the batched slots are forwarded to the main `DensityModel` +constructor and keep their meaning there. Since `ld` fills the gradient slot +with a callable, a *plain* backend in `hvp` would differentiate that callable — +which does not work for a gradient `ld` computes by AD, because its preparation +is tied to the input type it was made for and rejects the tangents an outer +pass pushes through. For an AD HVP here, pass a +`DifferentiationInterface.SecondOrder`: it differentiates the log-density twice +and so never goes near `ld`'s gradient. A hand-written `hvp` callable works too. + +The batched slots are how a LogDensityProblems model reaches the batched DEER +path. `ld` supplies no batched log-density, so `logdensity_batch` has to be +written out by hand, and the same `SecondOrder` caveat applies to `hvp_batch`. # Turing.jl / DynamicPPL example ```julia @@ -46,7 +58,14 @@ chain = sample(model, AdaptiveMALASampler(0.3; n_warmup=500), 2_000; If DynamicPPL is loaded, the simpler one-step constructor `DensityModel(mymodel(obs))` is also available and extracts parameter names automatically. """ -function ParallelMCMC.DensityModel(ld; param_names=nothing, hvp=nothing) +function ParallelMCMC.DensityModel( + ld; + param_names=nothing, + hvp=nothing, + logdensity_batch=nothing, + grad_logdensity_batch=nothing, + hvp_batch=nothing, +) caps = LogDensityProblems.capabilities(ld) caps isa LogDensityProblems.LogDensityOrder{0} && error( "LogDensityProblems model must support gradients (LogDensityOrder{1} or higher). " * @@ -58,7 +77,16 @@ function ParallelMCMC.DensityModel(ld; param_names=nothing, hvp=nothing) logp = ParallelMCMC.LogDensityProblemPrimal(ld) gradlogp = ParallelMCMC.LogDensityProblemGradient(ld) - return ParallelMCMC.DensityModel(logp, gradlogp, dim; param_names=param_names, hvp=hvp) + return ParallelMCMC.DensityModel( + logp, + gradlogp, + dim; + param_names=param_names, + hvp=hvp, + logdensity_batch=logdensity_batch, + grad_logdensity_batch=grad_logdensity_batch, + hvp_batch=hvp_batch, + ) end (l::ParallelMCMC.LogDensityProblemPrimal)(x) = LogDensityProblems.logdensity(l.ld, x) diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 2080015..ee1021f 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -177,13 +177,18 @@ end (c::_BatchHvpReverseClosure)(X, V) = pmcmc_dotsum(c.grad_batch(X), V) #= -Pick the AD-HVP fallback strategy from the user's backend. +Pick the AD-HVP fallback strategy from the user's backend. These two apply +when the HVP is one AD pass over a gradient we already have — a hand-written +`gradlogp`, which neither of them differentiates twice: ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Routes through the `pmcmc_matmul` frule. ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Routes through the matmul and dot/sum rrules. +An AD-derived gradient takes neither, going to `_make_hvp_fn_second_order` +below instead, so nothing here ever nests one DI call inside another. + These are singleton types rather than symbols so the choice dispatches statically — `_make_hvp_fn(_hvp_strategy(backend), ...)` resolves to one concrete method (and one concrete return type) at compile time, without @@ -304,6 +309,63 @@ function _make_hvp_batch_fn( return (X, V) -> _batch_hvp_via_grad_reverse_prepared(prep, X, V) end +#= +--------------------------------------------------------------------------- +True second-order HVP. `DI.hvp` with a `SecondOrder` backend differentiates +the log-density twice, so unlike the two strategies above this never touches +the gradient slot — the inner half of the pair is what produces the gradient. + +This is the path for a model whose gradient is itself AD-derived. Pushing +tangents through an already-prepared DI gradient would compute the same thing, +but the inner call falls out of its preparation as soon as the outer pass +hands it a tangent type the prep wasn't made for, so DI's own second-order +operator is both what the user asked for and the cheaper way to get it. + +Normalization applies to the outer half, the pass whose mode and annotations +the backend extensions care about. `_hvp_forward_backend` already unwraps a +`SecondOrder` to its outer, so it selects the right half out of `backend` +itself; the inner is left alone, since it is a plain first-order gradient of +the log-density and needs none of the pushforward-specific pinning. +--------------------------------------------------------------------------- +=# +function _normalized_second_order(backend::DI.SecondOrder) + return DI.SecondOrder(_hvp_forward_backend(backend), DI.inner(backend)) +end + +function _make_hvp_fn_second_order( + logdensity, backend::DI.SecondOrder, x_template::AbstractVector +) + so = _normalized_second_order(backend) + v_template = similar(x_template) + fill!(v_template, zero(eltype(x_template))) + prep = DI.prepare_hvp(logdensity, so, x_template, (v_template,)) + return function (pt, dir) + x_exec = _materialize_ad_vector(pt) + v_exec = _tangent_like(x_exec, dir) + return first(DI.hvp(logdensity, prep, so, x_exec, (v_exec,))) + end +end + +#= +Batched form. The scalar function is `sum(logdensity_batch(X))`, whose +Hessian is block-diagonal because the columns are independent, so its HVP +along `V` is the columnwise HVP — the same argument the batched gradient +rests on. +=# +function _make_hvp_batch_fn_second_order( + logdensity_batch_sum, backend::DI.SecondOrder, X_template::AbstractMatrix +) + so = _normalized_second_order(backend) + V_template = similar(X_template) + fill!(V_template, zero(eltype(X_template))) + prep = DI.prepare_hvp(logdensity_batch_sum, so, X_template, (V_template,)) + return function (X, V) + X_exec = _materialize_ad_matrix(X) + V_exec = _tangent_like(X_exec, V) + return first(DI.hvp(logdensity_batch_sum, prep, so, X_exec, (V_exec,))) + end +end + @inline function _rademacher!(z::AbstractArray{T}, rng::AbstractRNG) where {T} @inbounds for i in eachindex(z) z[i] = rand(rng, Bool) ? one(T) : -one(T) diff --git a/src/interface.jl b/src/interface.jl index 2c032f7..bbbcf7e 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -19,38 +19,43 @@ the log-density: DensityModel(logp, AutoForwardDiff(), dim) -A backend in `hvp` / `hvp_batch` differentiates whatever the gradient slot -holds; it does not take a second derivative of `logdensity`. Over an -AD-derived gradient that composition is second-order AD, and over a -hand-written gradient it is a single AD pass across your own code. Only the -outer half of a `DifferentiationInterface.SecondOrder` is used, the gradient -slot being the inner one. +How a backend in `hvp` / `hvp_batch` gets its second derivative depends on the +gradient slot. Over a hand-written gradient it is a single AD pass across your +own code. Over an AD-derived one it is true second-order AD, taken as +`DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` through DI's +own second-order operator. Passing a `SecondOrder` yourself always means the +latter — both passes are named, so it differentiates `logdensity` twice and the +gradient slot is not its inner pass even when hand-written. - `logdensity(x::AbstractVector) -> Real` - `grad_logdensity` — callable `x -> AbstractVector`, or a backend applied to `logdensity`. -- `hvp` — optional callable `(x, v) -> AbstractVector`, or a backend applied - to `grad_logdensity`. If `nothing`, DEER builds the HVP from the sampler's - `backend`. +- `hvp` — optional callable `(x, v) -> AbstractVector`, or a backend. If + `nothing`, DEER builds the HVP from the sampler's `backend`. - `logdensity_batch(X::AbstractMatrix) -> AbstractVector` — optional batched log-density over columns (callable only). Columns must be independent: element `t` of the result may depend on column `t` of `X` and nothing else. A batched gradient derived from this is one gradient of its sum, so coupling between columns would go unnoticed and give wrong derivatives. - `grad_logdensity_batch` — optional callable `X -> AbstractMatrix`, or a - backend applied to `logdensity_batch` (which must then be provided). Left - out alongside a `logdensity_batch`, it is derived from whichever backend is - available. -- `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend - applied to `grad_logdensity_batch`, whether that one was given or derived. + backend applied to `logdensity_batch`. Left out alongside a + `logdensity_batch`, it is derived when `grad_logdensity` is a backend. +- `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend, + resolved against `grad_logdensity_batch` the same way `hvp` is against + `grad_logdensity`. - `dim::Int` — dimensionality of the parameter space - `param_names` — optional collection of parameter names used in `FlexiChains` output. If `nothing` (the default), uses a single vector-valued parameter `:x` with shape `(dim,)`. See the [`Parameter names`](@ref parameter-names) section of the docs for more information. -When `logdensity_batch` and `grad_logdensity_batch` are provided, -`ParallelMALASampler` enables the batched DEER update path. +Both batched derivative slots require `logdensity_batch`, which the batched +update evaluates directly. `logdensity_batch` alone is allowed and is used to +score whole trajectories at once, without switching the batched update on. + +`ParallelMALASampler` enables the batched DEER update path once it has a +`logdensity_batch` and a batched gradient — `grad_logdensity_batch`, or one +derived from `logdensity_batch` when `grad_logdensity` is a backend. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel logdensity::F @@ -91,28 +96,22 @@ function DensityModel( "grad_logdensity must be a callable or an ADTypes.AbstractADType backend" ), ) - if grad_logdensity_batch isa AbstractADType && logdensity_batch === nothing - throw( - ArgumentError( - "grad_logdensity_batch given as an AD backend requires logdensity_batch to differentiate", - ), - ) - end - #= A batched gradient it can differentiate has to be reachable, but it need - not be in hand: `grad_logdensity_batch` can itself be derived from a - `logdensity_batch`, and whether a backend is around to do that isn't known - until the sampler shows up. So only the case with nothing batched at all is - rejected here; `_prepare_model` raises the rest. =# - if hvp_batch isa AbstractADType && - grad_logdensity_batch === nothing && - logdensity_batch === nothing - throw( - ArgumentError( - "hvp_batch given as an AD backend requires a batched gradient to " * - "differentiate: supply grad_logdensity_batch, or logdensity_batch for " * - "one to be derived from", - ), - ) + #= The batched DEER update evaluates `logdensity_batch` itself, so neither + batched derivative is usable without it — as a backend, because there would + be nothing to differentiate, and as a callable, because the path stays off + and the callable would never be reached. Rejecting both here rather than + silently ignoring them; a `logdensity_batch` on its own is fine, and + `_prepare_model` decides from the gradient slot whether the batched path + can actually run. =# + _batch_needs_logp(name) = throw( + ArgumentError( + "$name requires logdensity_batch: the batched DEER path evaluates the " * + "batched log-density, so it cannot run without one", + ), + ) + if logdensity_batch === nothing + grad_logdensity_batch === nothing || _batch_needs_logp("grad_logdensity_batch") + hvp_batch === nothing || _batch_needs_logp("hvp_batch") end return DensityModel( logdensity, @@ -162,8 +161,11 @@ _prepped_for(prepped::PreppedDensityModel, model::DensityModel) = prepped.source #= Resolved gradient wrappers. Structs rather than anonymous closures since DI keys preparations on function identity. `TX` is the input type the prep was -made for; anything else (e.g. the `Dual`s an outer AD pass pushes through -this gradient when forming an HVP) goes through unprepared `DI.gradient`. +made for; anything else falls back to unprepared `DI.gradient`, so a call with +an input the preparation does not cover still gives a right answer instead of +failing. Forming an HVP no longer takes that route — an AD-derived gradient +gets a true second-order operator rather than an outer pass pushed through +this one — so in a normal run the prepared branch is the one that fires. =# struct _ADGradient{F,B<:AbstractADType,P,TX} logdensity::F @@ -215,6 +217,60 @@ function _resolve_gradient_batch( ) end +#= +Resolve an HVP slot given as a backend. `grad` is the resolved gradient +callable and `grad_backend` the backend that produced it, or `nothing` when +the gradient slot held a callable. + +Three cases, and which one applies is decided by types alone so the branch +folds away and the returned closure type stays statically known: + + - A `SecondOrder` names both passes itself, so it always means a true + second-order derivative of the log-density. The gradient slot plays no + part in the HVP then (it is still the drift term the MALA step uses, just + not the HVP's inner pass) — including when it is a hand-written callable, + which a `SecondOrder` deliberately bypasses. + - A plain backend over an AD-derived gradient means the same thing: the pair + is `SecondOrder(hvp_backend, grad_backend)`, and handing it to DI as one + operator beats pushing tangents through the prepared gradient. + - A plain backend over a hand-written gradient is a single AD pass across + the user's own code, which is what the `HVPStrategy` paths do. +=# +function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) + if hvp_backend isa DI.SecondOrder + return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) + elseif grad_backend !== nothing + return DEER._make_hvp_fn_second_order( + logdensity, DI.SecondOrder(hvp_backend, grad_backend), x_template + ) + else + return DEER._make_hvp_fn( + DEER._hvp_strategy(hvp_backend), grad, hvp_backend, x_template + ) + end +end + +# Batched counterpart, on `sum(logdensity_batch(X))` for the second-order paths. +function _resolve_hvp_batch( + logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template +) + if hvp_backend isa DI.SecondOrder + return DEER._make_hvp_batch_fn_second_order( + _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template + ) + elseif grad_batch_backend !== nothing + return DEER._make_hvp_batch_fn_second_order( + _BatchLogdensitySum(logdensity_batch), + DI.SecondOrder(hvp_backend, grad_batch_backend), + X_template, + ) + else + return DEER._make_hvp_batch_fn( + DEER._hvp_strategy(hvp_backend), grad_batch, hvp_backend, X_template + ) + end +end + """ _prepare_model(model, x_template) -> PreppedDensityModel _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel @@ -224,11 +280,15 @@ at `x_template`. The two-argument form only does `grad_logdensity`, which is all the sequential samplers use, and leaves the DEER-only slots `nothing`. The four-argument form also does the HVP and batched slots, preparing those at -a `(dim, T)` template. A missing `hvp` / `hvp_batch` is derived from the -model's own `hvp` backend if it has one and the sampler's `backend` otherwise; -a missing `grad_logdensity_batch` likewise comes from the gradient slot's -backend or the sampler's. The batched slots are only filled when -`logdensity_batch` is there to derive them from. +a `(dim, T)` template. A missing `hvp` comes from the sampler's `backend`, and +a missing `hvp_batch` from the model's own `hvp` backend if it has one and the +sampler's otherwise. A missing `grad_logdensity_batch` is derived only when +`grad_logdensity` is a backend — never from the sampler's `backend`, which +would let it decide whether the batched update runs. + +The batched slots are filled only when `logdensity_batch` is present and a +batched gradient is reachable; otherwise the batched path stays off and the +unbatched update covers it. """ function _prepare_model(model::DensityModel, x_template::AbstractVector) grad = if model.grad_logdensity isa AbstractADType @@ -254,90 +314,96 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector) end function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, backend) - grad = if model.grad_logdensity isa AbstractADType - _resolve_gradient(model.logdensity, model.grad_logdensity, x_template) + grad_backend = + model.grad_logdensity isa AbstractADType ? model.grad_logdensity : nothing + grad = if grad_backend !== nothing + _resolve_gradient(model.logdensity, grad_backend, x_template) else model.grad_logdensity end - hvp = if model.hvp isa AbstractADType - b = model.hvp - DEER._make_hvp_fn(DEER._hvp_strategy(b), grad, b, x_template) - elseif model.hvp === nothing - backend === nothing && throw( + hvp = if model.hvp === nothing || model.hvp isa AbstractADType + hvp_backend = model.hvp === nothing ? backend : model.hvp + hvp_backend === nothing && throw( ArgumentError( "ParallelMALASampler needs a Hessian-vector product: supply `hvp` " * "on the DensityModel (callable or AD backend), or pass `backend=` " * "to ParallelMALASampler", ), ) - DEER._make_hvp_fn(DEER._hvp_strategy(backend), grad, backend, x_template) + _resolve_hvp(model.logdensity, grad, grad_backend, hvp_backend, x_template) else model.hvp end + #= A batched log-density with no batched gradient has one derived from the + model's own gradient backend, so that a log-density-only model reaches the + batched DEER path. The sampler's `backend` is deliberately not used here: + a model whose gradient slot is a backend has already opted into AD over its + own code, whereas one with a hand-written gradient has not, and letting the + HVP fallback double as a reason to run AD over `logdensity_batch` would + make `backend=` silently decide which update path runs. + + A `logdensity_batch` on its own is still useful without the batched path — + `_trajectory_logps` uses it — so failing to derive a batched gradient just + leaves the path off rather than raising. =# grad_batch = model.grad_logdensity_batch - hvp_batch = model.hvp_batch - - # Whatever can hand us a batched HVP: the model's own backend, else the sampler's. - hvp_source = model.hvp isa AbstractADType ? model.hvp : backend - - #= A batched log-density with no batched gradient still enables the batched - DEER path, deriving the gradient from the gradient slot's own backend if it - has one, else the sampler's. Only worth starting if the batched HVP can be - built too, and if either is unreachable the path just stays off: the - unbatched update is a complete fallback. =# - if model.logdensity_batch !== nothing && - grad_batch === nothing && - (hvp_batch !== nothing || hvp_source !== nothing) - grad_batch = if model.grad_logdensity isa AbstractADType - model.grad_logdensity - else - backend - end + if grad_batch === nothing && model.logdensity_batch !== nothing + grad_batch = grad_backend end + grad_batch_backend = grad_batch isa AbstractADType ? grad_batch : nothing + hvp_batch = model.hvp_batch batch_active = model.logdensity_batch !== nothing && grad_batch !== nothing - if grad_batch isa AbstractADType || - hvp_batch isa AbstractADType || - (batch_active && hvp_batch === nothing) + + if batch_active # Prepare on x0 in every column; zeros need not be in the support. X_template = similar(x_template, length(x_template), T) X_template .= x_template - if grad_batch isa AbstractADType + if grad_batch_backend !== nothing grad_batch = _resolve_gradient_batch( - model.logdensity_batch, grad_batch, X_template + model.logdensity_batch, grad_batch_backend, X_template ) end - if hvp_batch isa AbstractADType - #= Nothing derived a batched gradient above, so there is nothing for - this backend to differentiate. Unlike a missing `hvp_batch`, which - just leaves the batched path off, this was asked for explicitly. =# - grad_batch === nothing && throw( - ArgumentError( - "hvp_batch given as an AD backend has no batched gradient to " * - "differentiate: supply `grad_logdensity_batch` on the DensityModel, " * - "or a backend for one to be derived from `logdensity_batch` " * - "(`grad_logdensity`, or `backend=` on ParallelMALASampler)", - ), - ) - b = hvp_batch - hvp_batch = DEER._make_hvp_batch_fn( - DEER._hvp_strategy(b), grad_batch, b, X_template - ) - elseif batch_active && hvp_batch === nothing - hvp_source === nothing && throw( + + if hvp_batch === nothing || hvp_batch isa AbstractADType + # The model's own HVP backend if it has one, else the sampler's. + hvp_batch_backend = if hvp_batch === nothing + model.hvp isa AbstractADType ? model.hvp : backend + else + hvp_batch + end + hvp_batch_backend === nothing && throw( ArgumentError( "the batched DEER path needs a batched Hessian-vector product: " * "supply `hvp_batch` on the DensityModel (callable or AD backend), " * "or pass `backend=` to ParallelMALASampler", ), ) - hvp_batch = DEER._make_hvp_batch_fn( - DEER._hvp_strategy(hvp_source), grad_batch, hvp_source, X_template + hvp_batch = _resolve_hvp_batch( + model.logdensity_batch, + grad_batch, + grad_batch_backend, + hvp_batch_backend, + X_template, ) end + elseif hvp_batch !== nothing + #= The constructor rejects batched derivative slots without a + `logdensity_batch`, so this is a `logdensity_batch` plus an `hvp_batch` + with no batched gradient to pair it with, the gradient slot being a + callable there is nothing to derive one from. Raise rather than drop the + `hvp_batch`: it was supplied explicitly, and the alternative is silently + running the unbatched update. =# + throw( + ArgumentError( + "hvp_batch has no batched gradient to go with it: supply " * + "`grad_logdensity_batch` on the DensityModel (a callable, or a backend " * + "to derive one from `logdensity_batch`). A backend in `grad_logdensity` " * + "also derives one; `backend=` on ParallelMALASampler does not.", + ), + ) end return PreppedDensityModel( @@ -540,17 +606,11 @@ DEER-parallelized MALA sampler. Supported Jacobian modes are `:stoch_diag` (the default Hutchinson diagonal estimator) and `:diag` (exact diagonal via `D` JVPs). -`backend` is the fallback derivative source for what the `DensityModel` did not -bring: Hessian-vector products when it has no `hvp` / `hvp_batch` of its own, -and the batched gradient when it has a `logdensity_batch` but no -`grad_logdensity_batch` and its gradient slot is a callable rather than a -backend. A model that covers all of those does not need it. - -That second role means passing a `backend` can switch the batched DEER path on -for a model that would otherwise have run the unbatched update, which also puts -AD on the `logdensity_batch`. On GPU that brings the backend's restrictions -(see the GPU docs) to bear on a function nothing was differentiating before, so -supply `grad_logdensity_batch` if that matters. +`backend` is the fallback source of Hessian-vector products, used when the +`DensityModel` brings no `hvp` / `hvp_batch` of its own. That is all it does: it +never supplies a gradient, so it cannot change which update path runs or put AD +on a function the model did not already have a backend for. A model carrying its +own HVPs does not need it. """ struct ParallelMALASampler{FP<:AbstractFloat,CM,AD} <: AbstractMCMC.AbstractSampler epsilon::FP diff --git a/test/test-ADType-Slots.jl b/test/test-ADType-Slots.jl index afcc81c..fe7998e 100644 --- a/test/test-ADType-Slots.jl +++ b/test/test-ADType-Slots.jl @@ -19,6 +19,7 @@ gradlogp_batch_slots(X) = -X const D_SLOTS = 5 const CT_SLOTS = FlexiChains.FlexiChain{Symbol} +const DI_SLOTS = ParallelMCMC.DI @testset "constructor validation" begin # primal slots cannot be backends — nothing to derive them from @@ -26,13 +27,22 @@ const CT_SLOTS = FlexiChains.FlexiChain{Symbol} @test_throws ArgumentError DensityModel( logp_slots, gradlogp_slots, D_SLOTS; logdensity_batch=AutoForwardDiff() ) - # batched AD slots need their primal/gradient counterpart - @test_throws ArgumentError DensityModel( - logp_slots, gradlogp_slots, D_SLOTS; grad_logdensity_batch=AutoForwardDiff() - ) - @test_throws ArgumentError DensityModel( - logp_slots, gradlogp_slots, D_SLOTS; hvp_batch=AutoForwardDiff() - ) + #= Batched derivative slots need a logdensity_batch, callables included: the + batched update evaluates it, so without one the slot could never be reached. =# + for slot in (AutoForwardDiff(), gradlogp_batch_slots) + @test_throws ArgumentError DensityModel( + logp_slots, gradlogp_slots, D_SLOTS; grad_logdensity_batch=slot + ) + end + for slot in (AutoForwardDiff(), (X, V) -> -V) + @test_throws ArgumentError DensityModel( + logp_slots, gradlogp_slots, D_SLOTS; hvp_batch=slot + ) + end + # a logdensity_batch on its own is fine — it scores whole trajectories + @test DensityModel( + logp_slots, gradlogp_slots, D_SLOTS; logdensity_batch=logp_batch_slots + ) isa DensityModel # grad slot is mandatory (callable or backend) @test_throws ArgumentError DensityModel(logp_slots, nothing, D_SLOTS) @@ -162,50 +172,32 @@ end @test m_r.grad_logdensity_batch(X) ≈ -X @test m_r.hvp_batch(X, V) ≈ -V - # analytic gradient, so the sampler backend supplies it instead + #= A hand-written gradient is not opted into AD, so the sampler backend + does not derive a batched gradient from `logdensity_batch` and the + batched path stays off — the unbatched update covers it. Were it + otherwise, `backend=` would decide which update path runs. =# model_an = DensityModel( - logp_slots, gradlogp_slots, D_SLOTS; logdensity_batch=logp_batch_slots - ) - m_an = ParallelMCMC._prepare_model(model_an, x0, T, AutoForwardDiff()) - @test m_an.grad_logdensity_batch(X) ≈ -X - @test m_an.hvp_batch(X, V) ≈ -V - - #= No backend anywhere to derive from, so the batched path stays off - rather than erroring i.e., an analytic hvp still covers the unbatched one. =# - m_off = ParallelMCMC._prepare_model( - DensityModel( - logp_slots, - gradlogp_slots, - D_SLOTS; - hvp=(x, v) -> -v, - logdensity_batch=logp_batch_slots, - ), - x0, - T, - nothing, + logp_slots, + gradlogp_slots, + D_SLOTS; + hvp=(x, v) -> -v, # so the unbatched HVP is covered either way + logdensity_batch=logp_batch_slots, ) - @test m_off.grad_logdensity_batch === nothing - @test m_off.hvp_batch === nothing + for spl_backend in (AutoForwardDiff(), nothing) + m_an = ParallelMCMC._prepare_model(model_an, x0, T, spl_backend) + @test m_an.grad_logdensity_batch === nothing + @test m_an.hvp_batch === nothing + # still kept, since `_trajectory_logps` scores a whole block with it + @test m_an.logdensity_batch === logp_batch_slots + end end - @testset "hvp_batch backend with the batched gradient only derivable" begin + @testset "hvp_batch backend over a derived batched gradient" begin T = 8 X = randn(rng, D_SLOTS, T) V = randn(rng, D_SLOTS, T) - # no grad_logdensity_batch: derived from the sampler backend, then differentiated - model = DensityModel( - logp_slots, - gradlogp_slots, - D_SLOTS; - logdensity_batch=logp_batch_slots, - hvp_batch=AutoForwardDiff(), - ) - m_r = ParallelMCMC._prepare_model(model, x0, T, AutoForwardDiff()) - @test m_r.grad_logdensity_batch(X) ≈ -X - @test m_r.hvp_batch(X, V) ≈ -V - - # same, derived from the gradient slot's backend instead + # no grad_logdensity_batch: derived from the gradient slot's backend model_gs = DensityModel( logp_slots, AutoForwardDiff(), @@ -215,9 +207,12 @@ end hvp_batch=AutoForwardDiff(), ) m_gs = ParallelMCMC._prepare_model(model_gs, x0, T, nothing) + @test m_gs.grad_logdensity_batch(X) ≈ -X @test m_gs.hvp_batch(X, V) ≈ -V - #= Nowhere to derive the batched gradient from. =# + #= Hand-written gradient, so nothing derives a batched one, and an + `hvp_batch` was supplied explicitly: raise instead of silently running + the unbatched update. The sampler backend does not rescue this. =# model_nd = DensityModel( logp_slots, gradlogp_slots, @@ -226,14 +221,16 @@ end logdensity_batch=logp_batch_slots, hvp_batch=AutoForwardDiff(), ) - err = try - ParallelMCMC._prepare_model(model_nd, x0, T, nothing) - nothing - catch e - e + for spl_backend in (AutoForwardDiff(), nothing) + err = try + ParallelMCMC._prepare_model(model_nd, x0, T, spl_backend) + nothing + catch e + e + end + @test err isa ArgumentError + @test occursin("no batched gradient", err.msg) end - @test err isa ArgumentError - @test occursin("no batched gradient", err.msg) end @testset "model hvp backend feeds the batched HVP without a sampler backend" begin @@ -253,6 +250,138 @@ end end end +@testset "second-order HVP semantics" begin + #= How an `hvp` backend gets its second derivative depends on the gradient + slot, so the three cases are told apart with a gradient that is deliberately + not ∇logp: a path that differentiates the slot sees `-2x` and reports `-2v`, + one that differentiates `logdensity` twice reports `-v`. =# + off_grad(x) = -2 .* x + rng = MersenneTwister(61) + x0 = randn(rng, D_SLOTS) + v = randn(rng, D_SLOTS) + T = 8 + X = randn(rng, D_SLOTS, T) + V = randn(rng, D_SLOTS, T) + + @testset "plain backend over a hand-written gradient is one pass over it" begin + #= Both `HVPStrategy` paths stay reachable end to end and keep routing on + `DI.hvp_mode`: a hand-written gradient is differentiated once, never + twice, whichever direction the backend reports. `-2v` rather than `-v` is + what distinguishes that from the second-order paths below. =# + for backend in ( + AutoForwardDiff(), # ForwardOnGrad + AutoEnzyme(), # ForwardOnGrad + AutoEnzyme(; mode=Enzyme.Reverse), # ReverseOnGrad + ) + model = DensityModel(logp_slots, off_grad, D_SLOTS; hvp=backend) + m_r = ParallelMCMC._prepare_model(model, x0, T, nothing) + @test m_r.hvp(x0, v) ≈ -2 .* v + end + end + + @testset "an explicit SecondOrder differentiates logdensity twice" begin + #= Both passes are named, so the gradient slot is not the inner one even + when hand-written, and the inner half is honoured rather than dropped. =# + model = DensityModel( + logp_slots, + off_grad, + D_SLOTS; + hvp=DI_SLOTS.SecondOrder(AutoForwardDiff(), AutoForwardDiff()), + ) + m_r = ParallelMCMC._prepare_model(model, x0, T, nothing) + @test m_r.hvp(x0, v) ≈ -v + + # same when it arrives as the sampler's fallback backend + model_fb = DensityModel(logp_slots, off_grad, D_SLOTS) + m_fb = ParallelMCMC._prepare_model( + model_fb, x0, T, DI_SLOTS.SecondOrder(AutoForwardDiff(), AutoForwardDiff()) + ) + @test m_fb.hvp(x0, v) ≈ -v + end + + @testset "a backend over an AD-derived gradient composes into SecondOrder" begin + #= The gradient slot is a backend, so `hvp` becomes the outer half of a + true second-order pair rather than an AD pass over the prepared DI + gradient. Both readings agree numerically here (the derived gradient + really is ∇logp), so the check is structural: the slot has to hold the + closure the second-order factory builds, not the one a strategy builds. =# + second_order = ParallelMCMC.DEER._make_hvp_fn_second_order( + logp_slots, DI_SLOTS.SecondOrder(AutoForwardDiff(), AutoForwardDiff()), x0 + ) + + model = DensityModel(logp_slots, AutoForwardDiff(), D_SLOTS; hvp=AutoForwardDiff()) + m_r = ParallelMCMC._prepare_model(model, x0, T, nothing) + @test typeof(m_r.hvp) === typeof(second_order) + @test m_r.hvp(x0, v) ≈ -v + # the resolved gradient is still the prepared AD one + @test m_r.grad_logdensity isa ParallelMCMC._ADGradient + @test m_r.grad_logdensity(x0) ≈ -x0 + + # a hand-written gradient keeps the single-pass strategy closure instead + model_hand = DensityModel(logp_slots, off_grad, D_SLOTS; hvp=AutoForwardDiff()) + m_hand = ParallelMCMC._prepare_model(model_hand, x0, T, nothing) + @test typeof(m_hand.hvp) !== typeof(second_order) + end + + @testset "batched HVP follows the same three cases" begin + off_grad_batch(X) = -2 .* X + + model_hand = DensityModel( + logp_slots, + off_grad, + D_SLOTS; + hvp=(x, vv) -> -vv, + logdensity_batch=logp_batch_slots, + grad_logdensity_batch=off_grad_batch, + hvp_batch=AutoForwardDiff(), + ) + m_hand = ParallelMCMC._prepare_model(model_hand, x0, T, nothing) + @test m_hand.hvp_batch(X, V) ≈ -2 .* V + + model_so = DensityModel( + logp_slots, + off_grad, + D_SLOTS; + hvp=(x, vv) -> -vv, + logdensity_batch=logp_batch_slots, + grad_logdensity_batch=off_grad_batch, + hvp_batch=DI_SLOTS.SecondOrder(AutoForwardDiff(), AutoForwardDiff()), + ) + m_so = ParallelMCMC._prepare_model(model_so, x0, T, nothing) + @test m_so.hvp_batch(X, V) ≈ -V + + model_ad = DensityModel( + logp_slots, + AutoForwardDiff(), + D_SLOTS; + hvp=AutoForwardDiff(), + logdensity_batch=logp_batch_slots, + grad_logdensity_batch=AutoForwardDiff(), + hvp_batch=AutoForwardDiff(), + ) + m_ad = ParallelMCMC._prepare_model(model_ad, x0, T, nothing) + @test m_ad.hvp_batch(X, V) ≈ -V + end + + @testset "sampling with a SecondOrder hvp matches the analytic HVP" begin + model_so = DensityModel( + logp_slots, + gradlogp_slots, + D_SLOTS; + hvp=DI_SLOTS.SecondOrder(AutoForwardDiff(), AutoForwardDiff()), + ) + model_an = DensityModel(logp_slots, gradlogp_slots, D_SLOTS; hvp=(x, vv) -> -vv) + s = ParallelMALASampler(0.05; T=16) + c_so = sample( + MersenneTwister(62), model_so, s, 64; chain_type=CT_SLOTS, progress=false + ) + c_an = sample( + MersenneTwister(62), model_an, s, 64; chain_type=CT_SLOTS, progress=false + ) + @test c_so[:x] ≈ c_an[:x] + end +end + @testset "logp-only sampling matches analytic-gradient sampling" begin model_ad = DensityModel(logp_slots, AutoForwardDiff(), D_SLOTS) model_an = DensityModel(logp_slots, gradlogp_slots, D_SLOTS) diff --git a/test/test-Turing-Integration.jl b/test/test-Turing-Integration.jl index 03abb1a..9efe29b 100644 --- a/test/test-Turing-Integration.jl +++ b/test/test-Turing-Integration.jl @@ -132,6 +132,92 @@ end end end +@testset "DynamicPPLExt: SecondOrder hvp on a Turing model" begin + #= The gradient slot of a Turing model is DynamicPPL's own AD-prepared + gradient, and its preparation rejects the tangents an outer pass would push + through it — so a plain backend in `hvp` cannot differentiate it. A + `SecondOrder` differentiates the log-density twice instead, bypassing that + gradient, which is what makes an AD HVP reachable for a Turing model at all. + + normal_model(y) in unconstrained space is + logp(μ) = logpdf(Normal(0,1), μ) + logpdf(Normal(μ, 0.5), y), + so H = -1 - 1/0.5^2 = -5 and Hv = -5v. =# + so = ParallelMCMC.DI.SecondOrder(ADTypes.AutoForwardDiff(), ADTypes.AutoForwardDiff()) + model = DensityModel( + normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff(), hvp=so + ) + @test model.hvp === so + + prepped = ParallelMCMC._prepare_model(model, [0.0], 8, nothing) + @test prepped.hvp([0.0], [1.0]) ≈ [-5.0] + # the model brought its own HVP, so no sampler backend is needed + chain = sample( + MersenneTwister(12), + model, + ParallelMALASampler(0.02; T=8), + 64; + chain_type=VNChain, + progress=false, + ) + @test all(isfinite, vec(chain[@varname(μ)])) + + #= A plain backend is the case that cannot work. Preparing it succeeds — DI + only builds the pushforward against the Float64 template — and it is the + first call, pushing tangents into DynamicPPL's prepared gradient, that + fails. Pinned as a test so that if DynamicPPL ever lifts this, the + `SecondOrder`-only advice in the extension docstring gets revisited. =# + model_plain = DensityModel( + normal_model(TRUE_OBS); + ad_backend=ADTypes.AutoForwardDiff(), + hvp=ADTypes.AutoForwardDiff(), + ) + prepped_plain = ParallelMCMC._prepare_model(model_plain, [0.0], 8, nothing) + @test_throws Exception prepped_plain.hvp([0.0], [1.0]) +end + +@testset "DynamicPPLExt: batched slots reach the batched DEER path" begin + #= DynamicPPL supplies no batched log-density, so the batched slots are the + only way a Turing model reaches the batched update. Written out by hand for + normal_model, including the normalizing constants so that the log-densities + reported for a trajectory agree with `model.logdensity`. =# + σ = 0.5 + logp_b(X) = + vec(-0.5 .* X .^ 2 .- 0.5 .* ((TRUE_OBS .- X) ./ σ) .^ 2 .- log(2π) .- log(σ)) + grad_b(X) = -X .+ (TRUE_OBS .- X) ./ σ^2 + hvp_b(X, V) = (-1 - 1 / σ^2) .* V + + model = DensityModel( + normal_model(TRUE_OBS); + ad_backend=ADTypes.AutoForwardDiff(), + hvp=(x, v) -> (-1 - 1 / σ^2) .* v, + logdensity_batch=logp_b, + grad_logdensity_batch=grad_b, + hvp_batch=hvp_b, + ) + + @test model.logdensity_batch === logp_b + @test model.grad_logdensity_batch === grad_b + @test model.hvp_batch === hvp_b + + # the hand-written batched log-density agrees with the model's own, column by column + X = reshape([-0.5, 0.0, 0.7, 1.4], 1, 4) + @test logp_b(X) ≈ [model.logdensity(X[:, t]) for t in 1:size(X, 2)] + + prepped = ParallelMCMC._prepare_model(model, [0.0], 4, nothing) + @test prepped.grad_logdensity_batch === grad_b + @test prepped.hvp_batch === hvp_b + + chain = sample( + MersenneTwister(13), + model, + ParallelMALASampler(0.02; T=8), + 64; + chain_type=VNChain, + progress=false, + ) + @test all(isfinite, vec(chain[@varname(μ)])) +end + @testset "DynamicPPLExt: MvNormal(zeros(2), I) runs with ParallelMALA" begin model = DensityModel(mvnormal_2d_model(); ad_backend=ADTypes.AutoForwardDiff()) From 85212e64ab4fb8eaac4c56e0b6d7489a3d8a8402 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Tue, 28 Jul 2026 23:20:26 -0400 Subject: [PATCH 2/2] docs refresh --- docs/src/10-getting-started.md | 12 ++--- docs/src/15-gpu.md | 8 +-- docs/src/90-contributing.md | 2 +- ext/DynamicPPLExt.jl | 20 ++++---- ext/LogDensityProblemsExt.jl | 21 ++++---- src/DEER/DEER.jl | 39 +++++--------- src/interface.jl | 91 ++++++++++++--------------------- test/test-ADType-Slots.jl | 10 ++-- test/test-Turing-Integration.jl | 17 +++--- 9 files changed, 88 insertions(+), 132 deletions(-) diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index 516ad81..b24e81e 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -37,25 +37,23 @@ Any of the derivative slots (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, model = DensityModel(logp, AutoEnzyme(), 2; param_names=[:x1, :x2]) ``` -Backends are turned into prepared [DifferentiationInterface](https://github.com/JuliaDiff/DifferentiationInterface.jl) callables when sampling starts, and that preparation is reused for the rest of the run. Hand-written and AD-derived slots mix, so an analytical gradient with `hvp=AutoForwardDiff()` is fine. +Backends become prepared [DifferentiationInterface](https://github.com/JuliaDiff/DifferentiationInterface.jl) callables when sampling starts, and that preparation is reused for the rest of the run. Hand-written and AD-derived slots mix, so an analytical gradient with `hvp=AutoForwardDiff()` works. ### What a backend in `hvp` differentiates -How the HVP is taken depends on what the gradient slot holds: - | `grad_logdensity` | `hvp` | what runs | |---|---|---| | callable | backend | one AD pass over your gradient | | backend | backend | `SecondOrder(hvp, grad_logdensity)` on `logdensity` | | either | `SecondOrder(...)` | that pair on `logdensity`, gradient slot unused | -So a hand-written gradient with `hvp=AutoForwardDiff()` differentiates *your* code once, whereas a log-density-only model gets genuine second-order AD through DI's own second-order operator rather than one AD call nested inside another. Passing a `SecondOrder` yourself always means the log-density is differentiated twice — both passes are named, so the gradient slot is not the inner one even when you wrote it by hand. That is also the only AD route to an HVP for a Turing or LogDensityProblems model, whose gradient comes pre-prepared and cannot be differentiated again. +Naming both passes yourself is the one route that ignores the gradient slot, hand-written or not. It is also the only AD route to an HVP for a Turing or LogDensityProblems model, whose gradient arrives already prepared and cannot be differentiated again. -Which combinations actually work is up to the backends, and second-order support is the thinnest part of the ecosystem. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model; `AutoMooncake` does not, in either direction — it has no reverse-over-reverse, and its gradient will not accept an outer pass's tangents, so a `SecondOrder` with a forward outer fails too. Write `grad_logdensity` out by hand for Mooncake and the single pass over it works normally. None of the second-order combinations work on GPU yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37) and the [GPU page](15-gpu.md)). +Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents. Give Mooncake a hand-written `grad_logdensity` instead. No second-order pair works on GPU yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37) and the [GPU page](15-gpu.md)). -The batched pair works the same way, on `sum(logdensity_batch(X))`. Its gradient is the stacked per-column gradients only because the columns are independent, so `logdensity_batch` must not couple them. A `logdensity_batch` given without a `grad_logdensity_batch` has one derived when `grad_logdensity` is a backend; with a hand-written gradient the batched path stays off and the unbatched update covers it. Both batched derivative slots require `logdensity_batch`, but `logdensity_batch` on its own is fine — it is also used to score a whole trajectory at once. +The batched pair works the same way, on `sum(logdensity_batch(X))`. That sum's gradient is the stacked per-column gradients only because the columns are independent, so `logdensity_batch` must not couple them. Omitting `grad_logdensity_batch` derives one when `grad_logdensity` is a backend; with a hand-written gradient the batched path stays off and the unbatched update covers it. Both batched derivative slots require `logdensity_batch`, which is also useful on its own for scoring a whole trajectory at once. -`backend` on [`ParallelMALASampler`](@ref) is the fallback source of Hessian-vector products, for a model that brings no `hvp` / `hvp_batch` of its own. That is all it does: it never supplies a gradient, so it cannot change which update path runs. A model carrying its own HVPs can leave it out. +`backend` on [`ParallelMALASampler`](@ref) supplies Hessian-vector products for a model that brings no `hvp` / `hvp_batch` of its own, and nothing else. A model carrying its own can leave it out. --- diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 2a16cf3..ea99f64 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -233,10 +233,10 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You only supply `gradlogp` / `grad_logdensity_batch`.** The sampler builds the HVP by differentiating your gradient — either a forward-mode pushforward of `gradlogp` ([`ForwardOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), the default for most backends) or a reverse-mode gradient of `x -> dot(gradlogp(x), v)` ([`ReverseOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), used for `AutoMooncake` and `AutoZygote`). This is the **AD-HVP fallback**, and it is what the logistic-regression example above uses. !!! warning "Log-density-only models on GPU" - `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU. The HVP is then true second-order AD — `SecondOrder(hvp_backend, grad_backend)` on your log-density — which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it. The sequential samplers only need the gradient, so log-density-only models are fine there. + `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU. The HVP becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it. The sequential samplers only need the gradient, so log-density-only models work there. -!!! note "A hand-written gradient keeps AD out of the batched path" - The batched DEER path needs a batched gradient as well as a `logdensity_batch`. One is derived from `logdensity_batch` only when `grad_logdensity` is a backend — i.e. only for a model that already opted into AD over its own code. With a hand-written `grad_logdensity` and no `grad_logdensity_batch`, the batched path simply stays off and the unbatched update runs, so nothing you wrote by hand is ever handed to AD behind your back. The sampler's `backend` does not change this; it only ever supplies HVPs. +!!! note "A backend in `grad_logdensity` reaches `logdensity_batch` too" + The batched path needs a batched gradient, and derives one from `logdensity_batch` when `grad_logdensity` is a backend. That puts `logdensity_batch` under the same restrictions as the rest of your AD-visible code. Supply `grad_logdensity_batch` to avoid it. ### When the fallback is the right call @@ -249,7 +249,7 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **The HVP has a clean closed form.** Quadratic priors, Gaussian likelihoods, GLMs (logistic, Poisson, probit) — the second derivative is a known function of intermediate quantities you already compute in `gradlogp`. A few extra lines and you skip the AD pipeline entirely. - **Performance matters and the AD compile is heavy.** Enzyme and Mooncake both pay a one-shot compilation cost on the user's gradient. For long-running chains this amortizes, but for many short runs the analytical HVP wins. -- **You're hitting AD-backend-specific GPU restrictions.** The [Enzyme limitations](#2-enzyme-on-gpu-currently-needs-pmcmc_matmul-pmcmc_dot-pmcmc_dotsum) above (`pmcmc_*` wrappers, staged broadcasts) only matter when the AD backend is invoked. Supplying analytical HVP sidesteps them — your `gradlogp` and `hvp` can use plain `*`, `dot`, `sum`, and the sampler's `backend` can be omitted entirely because no AD is ever invoked. A model whose derivatives are all hand-written never reaches AD by any route, batched path included. +- **You're hitting AD-backend-specific GPU restrictions.** The [Enzyme limitations](#2-enzyme-on-gpu-currently-needs-pmcmc_matmul-pmcmc_dot-pmcmc_dotsum) above (`pmcmc_*` wrappers, staged broadcasts) only matter when the AD backend is invoked. Supplying analytical HVP sidesteps them — your `gradlogp` and `hvp` can use plain `*`, `dot`, `sum`, and the sampler's `backend` can be omitted entirely because no AD is ever invoked. - **You can reuse intermediates between gradient and HVP.** When `hvp` shares $X\beta$, $\sigma(X\beta)$, or similar with the gradient computation, an analytical version can be both faster *and* shorter than what AD produces. ### Same example with analytical HVP diff --git a/docs/src/90-contributing.md b/docs/src/90-contributing.md index 01e2934..0a22749 100644 --- a/docs/src/90-contributing.md +++ b/docs/src/90-contributing.md @@ -30,4 +30,4 @@ Before submitting a pull request, please make sure: - You have added tests for your feature or fix, where applicable. - You have formatted the code with `JuliaFormatter` using the version specified by the `test` project’s compatibility bounds. -- All tests pass. \ No newline at end of file +- All tests pass. diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index cabc650..064c6d6 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -20,17 +20,15 @@ triggers for this extension), plus any AD backend that is used. `ad_backend` is DynamicPPL's own `adtype`, not a `DensityModel` slot: it goes to the `LogDensityFunction` that fills the log-density and gradient slots, which is -why it takes a backend only and never a callable. Everything else here is a -`DensityModel` slot forwarded unchanged, so `hvp` and the batched slots take a -callable or a backend as they do on the main constructor. - -For `ParallelMALASampler`, which also needs an HVP, pass either a callable or a -`DifferentiationInterface.SecondOrder` — the latter differentiates the -log-density twice, bypassing DynamicPPL's gradient. A plain backend does not -work here: it would differentiate the gradient `ad_backend` produced, whose AD -preparation rejects the tangents an outer pass pushes through it. DynamicPPL -supplies no batched log-density either, so `logdensity_batch` has to be written -out by hand to reach the batched DEER path. +why it takes a backend only and never a callable. The rest are `DensityModel` +slots forwarded unchanged. + +`ParallelMALASampler` also needs an HVP. Give it a callable or a +`DifferentiationInterface.SecondOrder`, which differentiates the log-density and +so bypasses DynamicPPL's gradient. A plain backend fails, since it would +differentiate the gradient `ad_backend` produced and that preparation rejects an +outer pass's tangents. DynamicPPL supplies no batched log-density either, so +reaching the batched DEER path means writing `logdensity_batch` by hand. # Example ```julia diff --git a/ext/LogDensityProblemsExt.jl b/ext/LogDensityProblemsExt.jl index c940a5e..6f59578 100644 --- a/ext/LogDensityProblemsExt.jl +++ b/ext/LogDensityProblemsExt.jl @@ -20,18 +20,15 @@ The optional `param_names` keyword accepts a collection of parameter names that for the columns of the returned `FlexiChain` object. If omitted, a single vector-valued parameter named `:x` will be chosen, unless you also pass `param_names` to `sample(...)`. -`hvp` and the batched slots are forwarded to the main `DensityModel` -constructor and keep their meaning there. Since `ld` fills the gradient slot -with a callable, a *plain* backend in `hvp` would differentiate that callable — -which does not work for a gradient `ld` computes by AD, because its preparation -is tied to the input type it was made for and rejects the tangents an outer -pass pushes through. For an AD HVP here, pass a -`DifferentiationInterface.SecondOrder`: it differentiates the log-density twice -and so never goes near `ld`'s gradient. A hand-written `hvp` callable works too. - -The batched slots are how a LogDensityProblems model reaches the batched DEER -path. `ld` supplies no batched log-density, so `logdensity_batch` has to be -written out by hand, and the same `SecondOrder` caveat applies to `hvp_batch`. +`hvp` and the batched slots are forwarded to the main `DensityModel` constructor +and keep their meaning there, with one caveat. `ld` fills the gradient slot, and +a gradient `ld` computes by AD carries a preparation tied to its input type, so +it rejects the tangents a plain `hvp` backend would push through it. Use a +callable, or a `DifferentiationInterface.SecondOrder` which differentiates the +log-density instead. Same for `hvp_batch`. + +`ld` supplies no batched log-density, so reaching the batched DEER path means +writing `logdensity_batch` by hand. # Turing.jl / DynamicPPL example ```julia diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index ee1021f..03896ef 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -177,8 +177,8 @@ end (c::_BatchHvpReverseClosure)(X, V) = pmcmc_dotsum(c.grad_batch(X), V) #= -Pick the AD-HVP fallback strategy from the user's backend. These two apply -when the HVP is one AD pass over a gradient we already have — a hand-written +Pick the AD-HVP fallback strategy from the user's backend. These two apply when +the HVP is one AD pass over a gradient we already have, i.e., a hand-written `gradlogp`, which neither of them differentiates twice: ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Routes through the @@ -186,8 +186,7 @@ when the HVP is one AD pass over a gradient we already have — a hand-written ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Routes through the matmul and dot/sum rrules. -An AD-derived gradient takes neither, going to `_make_hvp_fn_second_order` -below instead, so nothing here ever nests one DI call inside another. +An AD-derived gradient takes neither and goes to `_make_hvp_fn_second_order`. These are singleton types rather than symbols so the choice dispatches statically — `_make_hvp_fn(_hvp_strategy(backend), ...)` resolves to one @@ -311,21 +310,17 @@ end #= --------------------------------------------------------------------------- -True second-order HVP. `DI.hvp` with a `SecondOrder` backend differentiates -the log-density twice, so unlike the two strategies above this never touches -the gradient slot — the inner half of the pair is what produces the gradient. - -This is the path for a model whose gradient is itself AD-derived. Pushing -tangents through an already-prepared DI gradient would compute the same thing, -but the inner call falls out of its preparation as soon as the outer pass -hands it a tangent type the prep wasn't made for, so DI's own second-order -operator is both what the user asked for and the cheaper way to get it. - -Normalization applies to the outer half, the pass whose mode and annotations -the backend extensions care about. `_hvp_forward_backend` already unwraps a -`SecondOrder` to its outer, so it selects the right half out of `backend` -itself; the inner is left alone, since it is a plain first-order gradient of -the log-density and needs none of the pushforward-specific pinning. +Second-order HVP, for a model whose gradient is itself AD-derived. `DI.hvp` +takes both passes over the log-density, so these never touch the gradient slot. +Preferred over pushing tangents through a prepared DI gradient, which drops out +of its preparation once the outer pass hands it an unexpected tangent type. + +Only the outer half is normalized: `_hvp_forward_backend` selects it out of the +pair. The inner is a plain first-order gradient and needs no pinning. + +The batched form differentiates `sum(logdensity_batch(X))`, whose Hessian is +block-diagonal by column independence, so its HVP along `V` is the columnwise +HVP. Same argument the batched gradient rests on. --------------------------------------------------------------------------- =# function _normalized_second_order(backend::DI.SecondOrder) @@ -346,12 +341,6 @@ function _make_hvp_fn_second_order( end end -#= -Batched form. The scalar function is `sum(logdensity_batch(X))`, whose -Hessian is block-diagonal because the columns are independent, so its HVP -along `V` is the columnwise HVP — the same argument the batched gradient -rests on. -=# function _make_hvp_batch_fn_second_order( logdensity_batch_sum, backend::DI.SecondOrder, X_template::AbstractMatrix ) diff --git a/src/interface.jl b/src/interface.jl index bbbcf7e..9568b0a 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -21,11 +21,10 @@ the log-density: How a backend in `hvp` / `hvp_batch` gets its second derivative depends on the gradient slot. Over a hand-written gradient it is a single AD pass across your -own code. Over an AD-derived one it is true second-order AD, taken as -`DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` through DI's -own second-order operator. Passing a `SecondOrder` yourself always means the -latter — both passes are named, so it differentiates `logdensity` twice and the -gradient slot is not its inner pass even when hand-written. +own code. Over an AD-derived one it is +`DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)`, taken through +DI's second-order operator. Passing a `SecondOrder` yourself always means the +latter, and bypasses the gradient slot even when you wrote it by hand. - `logdensity(x::AbstractVector) -> Real` - `grad_logdensity` — callable `x -> AbstractVector`, or a backend applied to @@ -50,11 +49,11 @@ gradient slot is not its inner pass even when hand-written. information. Both batched derivative slots require `logdensity_batch`, which the batched -update evaluates directly. `logdensity_batch` alone is allowed and is used to -score whole trajectories at once, without switching the batched update on. +update evaluates directly. `logdensity_batch` alone is allowed and scores whole +trajectories at once without switching the batched update on. -`ParallelMALASampler` enables the batched DEER update path once it has a -`logdensity_batch` and a batched gradient — `grad_logdensity_batch`, or one +`ParallelMALASampler` runs the batched DEER update once it has a +`logdensity_batch` and a batched gradient: either `grad_logdensity_batch`, or one derived from `logdensity_batch` when `grad_logdensity` is a backend. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel @@ -97,12 +96,10 @@ function DensityModel( ), ) #= The batched DEER update evaluates `logdensity_batch` itself, so neither - batched derivative is usable without it — as a backend, because there would - be nothing to differentiate, and as a callable, because the path stays off - and the callable would never be reached. Rejecting both here rather than - silently ignoring them; a `logdensity_batch` on its own is fine, and - `_prepare_model` decides from the gradient slot whether the batched path - can actually run. =# + batched derivative is usable without one: a backend would have nothing to + differentiate, and a callable would never be reached. Rejected here rather + than silently ignored. A `logdensity_batch` on its own is allowed, and + `_prepare_model` decides from the gradient slot whether the path can run. =# _batch_needs_logp(name) = throw( ArgumentError( "$name requires logdensity_batch: the batched DEER path evaluates the " * @@ -159,13 +156,10 @@ silently sample the wrong target. _prepped_for(prepped::PreppedDensityModel, model::DensityModel) = prepped.source === model #= -Resolved gradient wrappers. Structs rather than anonymous closures since DI -keys preparations on function identity. `TX` is the input type the prep was -made for; anything else falls back to unprepared `DI.gradient`, so a call with -an input the preparation does not cover still gives a right answer instead of -failing. Forming an HVP no longer takes that route — an AD-derived gradient -gets a true second-order operator rather than an outer pass pushed through -this one — so in a normal run the prepared branch is the one that fires. +Resolved gradient wrappers. Structs rather than anonymous closures since DI keys +preparations on function identity. `TX` is the input type the prep was made for; +anything else falls back to unprepared `DI.gradient` rather than failing. In a +normal run the prepared branch is the one that fires. =# struct _ADGradient{F,B<:AbstractADType,P,TX} logdensity::F @@ -218,23 +212,14 @@ function _resolve_gradient_batch( end #= -Resolve an HVP slot given as a backend. `grad` is the resolved gradient -callable and `grad_backend` the backend that produced it, or `nothing` when -the gradient slot held a callable. - -Three cases, and which one applies is decided by types alone so the branch -folds away and the returned closure type stays statically known: - - - A `SecondOrder` names both passes itself, so it always means a true - second-order derivative of the log-density. The gradient slot plays no - part in the HVP then (it is still the drift term the MALA step uses, just - not the HVP's inner pass) — including when it is a hand-written callable, - which a `SecondOrder` deliberately bypasses. - - A plain backend over an AD-derived gradient means the same thing: the pair - is `SecondOrder(hvp_backend, grad_backend)`, and handing it to DI as one - operator beats pushing tangents through the prepared gradient. - - A plain backend over a hand-written gradient is a single AD pass across - the user's own code, which is what the `HVPStrategy` paths do. +Resolve an HVP slot given as a backend. `grad_backend` is the backend that +produced `grad`, or nothing when the gradient slot held a callable. Dispatch is +on types alone, so the branch folds and the returned closure type stays +statically known. + +A `SecondOrder` bypasses the gradient slot even when that slot is hand-written: +naming both passes asks for two derivatives of `logdensity`. The slot is still +the drift term the MALA step uses. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) if hvp_backend isa DI.SecondOrder @@ -283,8 +268,8 @@ The four-argument form also does the HVP and batched slots, preparing those at a `(dim, T)` template. A missing `hvp` comes from the sampler's `backend`, and a missing `hvp_batch` from the model's own `hvp` backend if it has one and the sampler's otherwise. A missing `grad_logdensity_batch` is derived only when -`grad_logdensity` is a backend — never from the sampler's `backend`, which -would let it decide whether the batched update runs. +`grad_logdensity` is a backend, never from the sampler's `backend`, which would +let it decide whether the batched update runs. The batched slots are filled only when `logdensity_batch` is present and a batched gradient is reachable; otherwise the batched path stays off and the @@ -336,17 +321,11 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, model.hvp end - #= A batched log-density with no batched gradient has one derived from the - model's own gradient backend, so that a log-density-only model reaches the - batched DEER path. The sampler's `backend` is deliberately not used here: - a model whose gradient slot is a backend has already opted into AD over its - own code, whereas one with a hand-written gradient has not, and letting the - HVP fallback double as a reason to run AD over `logdensity_batch` would - make `backend=` silently decide which update path runs. - - A `logdensity_batch` on its own is still useful without the batched path — - `_trajectory_logps` uses it — so failing to derive a batched gradient just - leaves the path off rather than raising. =# + #= A batched log-density with no batched gradient gets one from the model's + own gradient backend, never the sampler's: a model with a hand-written + gradient has not opted into AD, and deriving one anyway would let `backend=` + decide which update path runs. Failing to derive leaves the path off rather + than raising, since `_trajectory_logps` uses `logdensity_batch` regardless. =# grad_batch = model.grad_logdensity_batch if grad_batch === nothing && model.logdensity_batch !== nothing grad_batch = grad_backend @@ -390,12 +369,8 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, ) end elseif hvp_batch !== nothing - #= The constructor rejects batched derivative slots without a - `logdensity_batch`, so this is a `logdensity_batch` plus an `hvp_batch` - with no batched gradient to pair it with, the gradient slot being a - callable there is nothing to derive one from. Raise rather than drop the - `hvp_batch`: it was supplied explicitly, and the alternative is silently - running the unbatched update. =# + #= Raise rather than drop it: `hvp_batch` was supplied explicitly, and + the alternative is silently running the unbatched update. =# throw( ArgumentError( "hvp_batch has no batched gradient to go with it: supply " * diff --git a/test/test-ADType-Slots.jl b/test/test-ADType-Slots.jl index fe7998e..bcae985 100644 --- a/test/test-ADType-Slots.jl +++ b/test/test-ADType-Slots.jl @@ -22,7 +22,7 @@ const CT_SLOTS = FlexiChains.FlexiChain{Symbol} const DI_SLOTS = ParallelMCMC.DI @testset "constructor validation" begin - # primal slots cannot be backends — nothing to derive them from + # primal slots cannot be backends: nothing to derive them from @test_throws ArgumentError DensityModel(AutoForwardDiff(), gradlogp_slots, D_SLOTS) @test_throws ArgumentError DensityModel( logp_slots, gradlogp_slots, D_SLOTS; logdensity_batch=AutoForwardDiff() @@ -39,7 +39,7 @@ const DI_SLOTS = ParallelMCMC.DI logp_slots, gradlogp_slots, D_SLOTS; hvp_batch=slot ) end - # a logdensity_batch on its own is fine — it scores whole trajectories + # a logdensity_batch on its own is allowed: it scores whole trajectories @test DensityModel( logp_slots, gradlogp_slots, D_SLOTS; logdensity_batch=logp_batch_slots ) isa DensityModel @@ -173,9 +173,9 @@ end @test m_r.hvp_batch(X, V) ≈ -V #= A hand-written gradient is not opted into AD, so the sampler backend - does not derive a batched gradient from `logdensity_batch` and the - batched path stays off — the unbatched update covers it. Were it - otherwise, `backend=` would decide which update path runs. =# + does not derive a batched gradient from `logdensity_batch` and the batched + path stays off. Were it otherwise, `backend=` would decide which update + path runs. =# model_an = DensityModel( logp_slots, gradlogp_slots, diff --git a/test/test-Turing-Integration.jl b/test/test-Turing-Integration.jl index 9efe29b..0eddef2 100644 --- a/test/test-Turing-Integration.jl +++ b/test/test-Turing-Integration.jl @@ -134,10 +134,10 @@ end @testset "DynamicPPLExt: SecondOrder hvp on a Turing model" begin #= The gradient slot of a Turing model is DynamicPPL's own AD-prepared - gradient, and its preparation rejects the tangents an outer pass would push - through it — so a plain backend in `hvp` cannot differentiate it. A - `SecondOrder` differentiates the log-density twice instead, bypassing that - gradient, which is what makes an AD HVP reachable for a Turing model at all. + gradient, whose preparation rejects the tangents an outer pass would push + through it, so a plain backend in `hvp` cannot differentiate it. A + `SecondOrder` differentiates the log-density instead, bypassing that gradient, + which is what makes an AD HVP reachable for a Turing model at all. normal_model(y) in unconstrained space is logp(μ) = logpdf(Normal(0,1), μ) + logpdf(Normal(μ, 0.5), y), @@ -161,11 +161,10 @@ end ) @test all(isfinite, vec(chain[@varname(μ)])) - #= A plain backend is the case that cannot work. Preparing it succeeds — DI - only builds the pushforward against the Float64 template — and it is the - first call, pushing tangents into DynamicPPL's prepared gradient, that - fails. Pinned as a test so that if DynamicPPL ever lifts this, the - `SecondOrder`-only advice in the extension docstring gets revisited. =# + #= Preparing a plain backend succeeds, since DI only builds the pushforward + against the Float64 template. The first call is what fails, pushing tangents + into DynamicPPL's prepared gradient. Pinned so that if DynamicPPL ever lifts + this, the `SecondOrder`-only advice in the extension docstring is revisited. =# model_plain = DensityModel( normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff(),