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..b24e81e 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -37,11 +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. -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. +| `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 | + +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 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))`. 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) 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 9153303..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 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 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 `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 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. 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. - **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 0e5e92a..064c6d6 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,18 @@ 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. 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 using Turing, ParallelMCMC, FlexiChains, ForwardDiff @@ -31,7 +44,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 +60,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..6f59578 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,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(...)`. -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, 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 @@ -46,7 +55,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 +74,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..03896ef 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -177,13 +177,17 @@ 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, i.e., 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 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 concrete method (and one concrete return type) at compile time, without @@ -304,6 +308,53 @@ function _make_hvp_batch_fn( return (X, V) -> _batch_hvp_via_grad_reverse_prepared(prep, X, V) end +#= +--------------------------------------------------------------------------- +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) + 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 + +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..9568b0a 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -19,38 +19,42 @@ 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 +`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 `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 scores whole +trajectories at once without switching the batched update on. + +`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 logdensity::F @@ -91,28 +95,20 @@ 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 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 " * + "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, @@ -160,10 +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 (e.g. the `Dual`s an outer AD pass pushes through -this gradient when forming an HVP) goes through unprepared `DI.gradient`. +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 @@ -215,6 +211,51 @@ function _resolve_gradient_batch( ) end +#= +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 + 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 +265,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 +299,86 @@ 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 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 - 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 + #= 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 " * + "`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 +581,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..bcae985 100644 --- a/test/test-ADType-Slots.jl +++ b/test/test-ADType-Slots.jl @@ -19,20 +19,30 @@ 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 + # 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() ) - # 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 allowed: 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. 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..0eddef2 100644 --- a/test/test-Turing-Integration.jl +++ b/test/test-Turing-Integration.jl @@ -132,6 +132,91 @@ 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, 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), + 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(μ)])) + + #= 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(), + 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())