Cleans up some docs prose and adds more tests

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
 

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.
 
 ---
 

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

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.
+- All tests pass.

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
 
 ######################

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)