diff --git a/CHANGELOG.md b/CHANGELOG.md index d2d204a..44c6194 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,107 +10,123 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - The derivative slots of `DensityModel` (`grad_logdensity`, `hvp`, - `grad_logdensity_batch`, `hvp_batch`) now take an `ADTypes.AbstractADType` - in place of a callable, so `DensityModel(logp, AutoForwardDiff(), dim)` - builds a model from the log-density alone. Backends are turned into - prepared DifferentiationInterface callables when sampling starts, and that - preparation is reused across steps (#40, #52). The prepared model rides - along in the sampler state to get that reuse; a state handed to `step` for a - different model, which `initial_state` allows, is re-prepared from the model - passed rather than reused, so the model given to `sample` is the one sampled. -- `ParallelMALASampler`'s `backend` keyword is now optional. It is only the - fallback source of Hessian-vector products, so a `DensityModel` carrying - 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. + `grad_logdensity_batch`, `hvp_batch`) now accept an `ADTypes.AbstractADType` + in place of a callable, so `DensityModel(logp, AutoForwardDiff(), dim)` builds + a model from the log-density alone (#40, #52). Backends become prepared + DifferentiationInterface callables when sampling starts, and the prepared + model rides along in the sampler state so the preparation is reused. A state + handed to `step` for a different model — which `initial_state` allows — is + re-prepared, so the model given to `sample` is the one sampled. +- `ParallelMALASampler`'s `backend` keyword is now optional: it is only the + fallback source of Hessian-vector products, so a `DensityModel` carrying its + own `hvp` / `hvp_batch` does not need one. Without it, a batched HVP comes + from the model's own `hvp` backend (#52). - A `logdensity_batch` given without a `grad_logdensity_batch` now has the - 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. + batched gradient derived from it when `grad_logdensity` is a backend, instead + of leaving the batched DEER path switched off. `hvp_batch` can be a backend in + that case too, and differentiates the derived gradient (#52). +- An HVP backend over an AD-derived gradient is now true second-order AD, + `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed to + `DI.hvp`, rather than an outer AD pass over the prepared DI gradient — which + dropped out of its preparation as soon as tangents were pushed through it + (#37). Around 10x fewer allocations for a `logdensity`-only model. A backend + over a hand-written gradient still differentiates that gradient once. +- `hvp` / `hvp_batch` accept a `SecondOrder` with both halves honoured, so the + log-density is differentiated twice and the gradient slot is not the inner + pass. The inner half used to be discarded. This is the one AD route to an HVP + for a Turing or LogDensityProblems model, whose gradient arrives already + prepared and cannot be differentiated again. +- New `ReactantExt`. `ADTypes.AutoReactant()` in a derivative slot, or as the + sampler `backend`, traces the derivative with Enzyme-MLIR and compiles it to + an XLA executable via Reactant.jl, off Enzyme's LLVM pipeline and off + DifferentiationInterface entirely, which yields a true second-order HVP for + a log-density-only model (#37, #52). Requires `using Reactant` and a + Reactant-traceable log-density. `AutoReactant` does not pair with a + DifferentiationInterface backend, a `LogDensityProblems` gradient, or a + `SecondOrder` across the two passes of an HVP, and doing so raises an + `ArgumentError` before any AD runs, as does a non-default + `AutoReactant(; mode=...)`. Two caveats, spelled out in `ext/ReactantExt.jl`'s + module docstring and the Reactant section of `docs/src/15-gpu.md`: a traced + log-density must be pure with respect to the data it captures, since a captured + array mutated after preparation stays frozen at its old value in every + derivative compiled from it, and the compiled derivative runs on whatever + device Reactant's XLA client targets, which need not be the GPU the model's + arrays live on. - 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 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 helper it still uses agrees: normalization - applies to the outer pass. Unwrapping to that half happens before the - normalization hook is dispatched on, so a `SecondOrder(AutoEnzyme(), ...)` still - reaches `EnzymeExt` and gets its function annotation filled in rather than - running as a bare `AutoEnzyme()`. -- Backend normalization no longer picks a differentiation mode on the user's - behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward` (with - `set_runtime_activity`) onto an `AutoEnzyme()` left mode-agnostic, on the - grounds that reverse mode hit a gc-transition abort on GPU and that composed - `pmcmc_matmul` calls needed runtime activity. The `pmcmc_*` Enzyme rules keep - Enzyme off both paths on their own now, so the pin bought nothing — and it cost - correctness, because it silently rewrote the direction of a `SecondOrder`'s - outer half. `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is - reverse-over-forward to `hvp_mode`, its inner half being forward-only, and came - out forward-over-forward. Normalization now fills in only - `function_annotation=Enzyme.Const`, which is about this package's own read-only - HVP wrappers rather than about Enzyme's mode, and leaves `mode` exactly as given - — unset included, for DI to resolve from the operator it runs. `hvp_mode` is - therefore identical before and after normalization for every backend pair. - - A mode set explicitly was never overridden, so only mode-agnostic backends were - affected, and the HVP was a correct HVP either way; what changes is that the - composition asked for is the one that runs. The two normalization hooks - (`_hvp_forward_backend`, `_hvp_closure_backend`) collapse into a single - `_normalized_backend`, since without a mode to choose they no longer differ. - Users relying on a plain `AutoEnzyme()` being run forward should now pass - `AutoEnzyme(; mode=Enzyme.Forward)` explicitly. - ### Changed - The AD-HVP fallback strategy (forward-on-grad vs reverse-on-grad) now comes from DifferentiationInterface's `hvp_mode` trait rather than a hardcoded per-backend list, so `AutoEnzyme(mode=Enzyme.Reverse)` routes to the 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 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. +- A model whose `grad_logdensity` is a backend now runs the batched update where + it used to run the unbatched one, since a `logdensity_batch` without a + `grad_logdensity_batch` has one derived for it. On GPU that puts the backend's + restrictions (`pmcmc_*` wrappers for Enzyme) on a `logdensity_batch` nothing + was differentiating before; supply `grad_logdensity_batch` to keep AD out of + it. A hand-written `grad_logdensity` is unaffected. - `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. + for a model with a hand-written gradient, which let a keyword that reads as an + HVP fallback decide which update path ran. Pass `grad_logdensity_batch` + explicitly, or a backend in `grad_logdensity` to derive one 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` on its own used to be accepted + and then ignored. `logdensity_batch` alone still scores 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. + now raises rather than falling back to the unbatched update. + +### Fixed + +- The reverse-on-grad HVP path differentiated with `DI.inner(backend)` while 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 second-order path instead of either strategy, and + normalization applies to the outer half after unwrapping, so + `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` rather than running + as a bare `AutoEnzyme()`. +- Backend normalization no longer picks a differentiation mode on the user's + behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward`, with + `set_runtime_activity`, onto a mode-agnostic `AutoEnzyme()`, against a + gc-transition abort on GPU and an `EnzymeRuntimeActivityError` on composed + `pmcmc_matmul` calls. The `pmcmc_*` Enzyme rules keep Enzyme off both paths on + their own, so the pin bought nothing and cost correctness: it rewrote the + direction of a `SecondOrder`'s outer half, turning + `SecondOrder(AutoEnzyme(), AutoForwardDiff())` — reverse-over-forward to + `hvp_mode`, its inner half being forward-only — into forward-over-forward. + Normalization now fills in `function_annotation=Enzyme.Const` and leaves `mode` + exactly as given, unset included, so `hvp_mode` reads the same before and after + it for every pair. Only mode-agnostic backends were affected and the HVP was + correct either way; what changes is that the composition asked for is the one + that runs. The two hooks `_hvp_forward_backend` and `_hvp_closure_backend` + collapse into one `_normalized_backend`. Pass + `AutoEnzyme(; mode=Enzyme.Forward)` to keep the old direction. +- `_check_reactant_pair` now runs before `_prepare_model` resolves a gradient, so + a mismatched `AutoReactant` pairing is reported immediately instead of after a + full XLA compile, and `AutoReactant` nested inside a `SecondOrder` is rejected + there too rather than reaching `DI.prepare_hvp`. An `AutoReactant` gradient + composes through the same second-order branch as any other AD-derived gradient, + and `_hvp_strategy(::AutoReactant)` is now live rather than a dead branch. +- `test/test-Reactant-HVP.jl` asserts a posterior mean against a target with a + known mean, and cross-checks against the analytic HVP on the same noise tape. + `size(chain)` and `all(isfinite, ...)` pass even for a badly wrong HVP, since + DEER's Newton iteration then just fails to converge rather than producing + `NaN`s. ### Removed - `DynamicPPLExt` no longer requires `ForwardDiff` as a triggering library to load. +- `Reactant` moved from `test/Project.toml`'s `[deps]` to `[extras]`, and + `test/test-Reactant-HVP.jl` is skipped unless `PARALLELMCMC_TEST_REACTANT` is + set. `Reactant_jll` ships a prebuilt XLA and is a large download most CI runs + and most local `Pkg.test()` calls should not have to pay for. Opt in by setting + the env var and adding `Reactant` to the test environment. ## [0.2.0] - 2026-06-29 diff --git a/Project.toml b/Project.toml index 82d250a..40b34e9 100644 --- a/Project.toml +++ b/Project.toml @@ -20,12 +20,14 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] DynamicPPLExt = ["DynamicPPL", "FlexiChains", "LogDensityProblems"] EnzymeExt = "Enzyme" LogDensityProblemsExt = "LogDensityProblems" +ReactantExt = ["Reactant", "Enzyme"] [compat] ADTypes = "1.21.0" @@ -42,6 +44,7 @@ LogDensityProblems = "2" Mooncake = "0.5.26" OrderedCollections = "1" Random = "1" +Reactant = "0.2.278" Statistics = "1" Zygote = "0.7.10" julia = "1.10" diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index b24e81e..50f4d4a 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -49,11 +49,11 @@ Backends become prepared [DifferentiationInterface](https://github.com/JuliaDiff 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)). +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, so give Mooncake a hand-written `grad_logdensity`. On GPU no DI-driven second-order pair works yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). `AutoReactant()` in both slots is the exception: it skips DI and traces forward-over-reverse straight from the log-density as one compiled XLA program. See 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. +`backend` on [`ParallelMALASampler`](@ref) supplies Hessian-vector products for a model that brings no `hvp` / `hvp_batch` of its own, and nothing else, so a model carrying its own can leave it out. --- diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index ea99f64..29aab89 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -232,12 +232,36 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You supply `hvp` / `hvp_batch`.** These run as plain kernels. The AD backend is never invoked for HVPs. - **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 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 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. +### Reactant HVPs, off the DI path + +`ADTypes.AutoReactant()` (requires `using Reactant`) goes another way. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), never touching Enzyme's LLVM pipeline, so neither the gc-transition abort nor the `pmcmc_*` wrappers above apply. It gives a true second-order HVP for a log-density-only model: + +```julia +using Reactant, ADTypes + +model = DensityModel(logp, AutoReactant(), D) # no hand-written gradient +sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) +``` + +With both the gradient slot and the HVP source `AutoReactant()`, the HVP compiles as explicit forward-over-reverse from the raw log-density, as its own XLA program. That is a separate executable from the plain gradient's: what gets fused is the HVP's two AD passes, not the gradient and the HVP. With a hand-written `gradlogp` it compiles a forward pushforward of your gradient instead. + +!!! warning "Two silent failure modes" + **Captured data is frozen at compile time.** `Reactant.@compile` bakes any plain `Array` or `Ref` reached through the closure into the executable as a constant. A `logdensity` written as `x -> f(x, data)` whose `data` you later mutate (`data .= new_values`) keeps returning the pre-mutation derivative from every executable compiled before the mutation, with no error and no warning, which biases the samples. Traced functions must be pure with respect to what they capture; pass data that can change in as an argument. + + **"GPU" here describes your array type, not the device Reactant runs on.** Which XLA client Reactant compiles for is a Reactant/`Reactant_jll`-wide setting (`Reactant.set_default_backend`) that nothing in this package controls, and it is `"cpu"` unless a GPU client was selected explicitly. Preparing a model whose `x_template` is a `CuArray` against a CPU client still compiles and still gives correct answers, but every call round-trips `CuArray → host → XLA-CPU → host → CuArray`. `_prepare_model` warns on that combination and can do nothing else about it; point Reactant at a GPU client yourself. + +Other caveats: + +- The traced function must be **Reactant-traceable**: plain array operations. DynamicPPL-built log-densities do not trace as-is. +- **Reactant does not mix with the DI backends across the two passes of an HVP.** Both `grad_logdensity` and the HVP source take `AutoReactant()`, or neither does; a hand-written gradient pairs with either. Mixing raises an `ArgumentError` at preparation time, since Reactant cannot trace a DI-prepared gradient and DI cannot differentiate a compiled XLA executable. Same for a `LogDensityProblems`/Turing gradient, and for nesting `AutoReactant()` inside a `DifferentiationInterface.SecondOrder`. +- **`AutoReactant(; mode=...)` raises** rather than being ignored. Derivatives always trace as Enzyme reverse-mode with the HVP as forward-over-that, whatever `mode` says. +- **Every call crosses a marshalling boundary** between package arrays and Reactant's XLA device memory: two host round-trips and a handful of allocations, not a fused in-place path. Executables are shape-specialized at preparation time, and `@compile` does not memoize across preparations, so every `sample()` call recompiles every `AutoReactant` slot rather than only the first one in the process. + +DifferentiationInterface cannot interoperate with Reactant. When that support lands ([DI#918](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/918)) this path folds into the ordinary `hvp_mode` routing with no user-facing change. + ### When the fallback is the right call - **Complex or composed models.** Bayesian neural nets, hierarchical models with many transformations, mixtures, or anything where the Hessian has no convenient closed form. Deriving and maintaining `hvp` by hand for these is error-prone; AD removes a whole class of bugs. diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index b7f2fb9..aac80ed 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -6,6 +6,15 @@ CurrentModule = ParallelMCMC This page documents all public types and functions exported by ParallelMCMC.jl. +## Reactant + +Loading `Reactant` (`using Reactant`) enables `ADTypes.AutoReactant()` as a +derivative-slot backend on `DensityModel` and as `ParallelMALASampler`'s +`backend`. It is not a `DifferentiationInterface` backend like the others. +[GPU Execution](15-gpu.md) covers what it does, its pairing rules, and its +caveats, in particular that traced log-densities must be pure with respect +to any data they capture. + ## Extension constructors `DensityModel` also has extension constructors for common probabilistic-programming interfaces: diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index 064c6d6..1df483b 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -18,17 +18,13 @@ 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. +`ad_backend` is DynamicPPL's own `adtype` rather than a `DensityModel` slot: it +goes to the `LogDensityFunction` that fills the log-density and gradient slots, +which is why it takes a backend and never a callable. The remaining keywords are +`DensityModel` slots, forwarded unchanged — see the `LogDensityProblemsExt` +constructor's docstring for the `hvp`/`hvp_batch` caveats that come with an +AD-derived gradient (no plain backend, no `AutoReactant()`, no batched +log-density without writing one by hand). # Example ```julia diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 82e63f8..33899fd 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -23,19 +23,14 @@ using Enzyme.EnzymeCore.EnzymeRules: Normalization of the user's `AutoEnzyme` for DEER's AD-HVP paths: fill in `function_annotation=Enzyme.Const` when they left it open, so Enzyme doesn't throw `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure` wrappers, which capture `gradlogp`. Those wrapper types -belong to this package, so declaring them constant is this package's business. +`_BatchHvpReverseClosure` wrappers that capture `gradlogp`. -`mode` is passed through exactly as given, unset included. Choosing a direction -on the user's behalf is not our call: a mode they set is a decision, and an unset -one is DI's to resolve from the operator it runs. - -An earlier version pinned `mode=Enzyme.Forward` here (with -`set_runtime_activity`) against the gc-transition abort on GPU and -`EnzymeRuntimeActivityError` on composed `pmcmc_matmul` calls. The rules below -keep Enzyme off both paths on their own, so the pin bought nothing and cost -correctness: it silently rewrote the direction of a `SecondOrder`'s outer half -(see `DEER._normalized_backend`). +`mode` passes through exactly as given, unset included. An earlier version pinned +`mode=Enzyme.Forward` here, with `set_runtime_activity`, against the GPU +gc-transition abort and the `EnzymeRuntimeActivityError` on composed +`pmcmc_matmul` calls. The rules below keep Enzyme off both paths by themselves, +so the pin bought nothing and cost correctness: it rewrote the direction of a +`SecondOrder`'s outer half (see `DEER._normalized_backend`). =# function DEER._normalized_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} A === Nothing || return backend diff --git a/ext/LogDensityProblemsExt.jl b/ext/LogDensityProblemsExt.jl index 6f59578..b146b2b 100644 --- a/ext/LogDensityProblemsExt.jl +++ b/ext/LogDensityProblemsExt.jl @@ -24,8 +24,10 @@ parameter named `:x` will be chosen, unless you also pass `param_names` to `samp 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`. +callable, or a `DifferentiationInterface.SecondOrder`, which differentiates the +log-density instead. Same for `hvp_batch`. `ADTypes.AutoReactant()` is out here +too, DI or no DI: Reactant cannot trace the DynamicPPL/LogDensityProblems +machinery `ld`'s gradient dispatches into. `ld` supplies no batched log-density, so reaching the batched DEER path means writing `logdensity_batch` by hand. diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl new file mode 100644 index 0000000..116b0dd --- /dev/null +++ b/ext/ReactantExt.jl @@ -0,0 +1,196 @@ +module ReactantExt + +#= +Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in any +derivative slot of `DensityModel`, or as the sampler `backend`. + +Two silent failure modes: + - Captured data is frozen at compile time. `@compile` bakes any plain + `Array`/`Ref` reached through the closure in as a constant, so mutating it + later keeps returning the pre-mutation derivative with no error or + warning. Traced functions must be pure w.r.t. their captures; pass mutable + data in as an argument instead. + - HVPs are explicit forward-over-(gradient) compositions, NEVER `Enzyme.hvp`, + which silently returns zeros under `@compile`. + +Other limitations: + - The traced function must be Reactant-traceable (plain array ops). + DynamicPPL-built log-densities do NOT trace as-is. + - Reactant's XLA client (`Reactant.set_default_backend`) is process-wide and + independent of where package arrays live. Preparing on a `CuArray` while it + targets `"cpu"` still compiles and is still correct, but every call + round-trips through the host; `_warn_reactant_host_roundtrip` warns once + and can't fix it. + - Executables are shape-specialized to the preparation templates and are not + memoized across preparations, so every `sample()` call recompiles every + `AutoReactant` slot the model uses. +=# + +using ParallelMCMC: ParallelMCMC +using ParallelMCMC.DEER: DEER +using ADTypes: ADTypes, AutoReactant, AutoEnzyme +using Reactant: Reactant, @compile +using Enzyme: Enzyme +using CUDA: CUDA + +#= +`AutoReactant()` defaults to `AutoReactant(; mode=AutoEnzyme())`, i.e. +`AutoReactant{AutoEnzyme{Nothing,Nothing}}` — the exact type matched below. +Anything else names an Enzyme mode or annotation, and this extension honours +neither. +=# +_check_reactant_mode(::AutoReactant{AutoEnzyme{Nothing,Nothing}}) = nothing +function _check_reactant_mode(backend::AutoReactant) + return throw( + ArgumentError( + "AutoReactant(; mode=$(backend.mode)) is not supported: gradients always " * + "trace as Enzyme reverse-mode and HVPs as forward-over-that, whatever " * + "`mode` says. Use the default AutoReactant().", + ), + ) +end + +#= +Warn once per compiled slot when the template is not a plain `Array` (so looks +like it lives on a GPU) while Reactant's default XLA client targets "cpu". +=# +function _reactant_client_platform() + return try + string(Reactant.XLA.platform_name(Reactant.XLA.default_backend())) + catch + nothing + end +end + +# Only `CuArray` implies a device round-trip; other non-`Array` templates (e.g. +# `SubArray` views) already live on the host and would be false positives. +_warn_reactant_host_roundtrip(::AbstractArray) = nothing +function _warn_reactant_host_roundtrip(x::CUDA.CuArray) + if _reactant_client_platform() == "cpu" + @warn "AutoReactant: preparing on a $(typeof(x)), but Reactant's default XLA " * + "client targets \"cpu\". Every call will round-trip to the host and back " * + "instead of running where the array lives, which is likely slower than not " * + "using Reactant at all. Point Reactant at a GPU client with " * + "`Reactant.set_default_backend(\"gpu\")` if one is available." maxlog = 1 + end + return nothing +end + +# Host-materialize for the Reactant boundary: Reactant manages its own (XLA) +# device memory, so we round-trip through a plain host Array regardless of +# where the package's arrays live (Vector / CuArray / SubArray). +_host(x::Array) = x +_host(x::AbstractArray) = Array(x) + +#= +Eltype comes from `out`, what Reactant actually computed, not from `template`: a +traced computation that promotes internally would otherwise be narrowed back to +the template's eltype on the way out. +=# +function _from_host(template::AbstractArray, out) + out_h = Array(out) + res = similar(template, eltype(out_h), size(out_h)) + copyto!(res, out_h) + return res +end + +#= +Compile `core` for the template shapes once and return a closure that +marshals package arrays <-> Reactant arrays. +=# +function _compiled(core, t1::AbstractArray) + _warn_reactant_host_roundtrip(t1) + r1 = Reactant.to_rarray(_host(t1)) + compiled = @compile core(r1) + return x -> _from_host(x, compiled(Reactant.to_rarray(_host(x)))) +end + +function _compiled(core, t1::AbstractArray, t2::AbstractArray) + _warn_reactant_host_roundtrip(t1) + r1 = Reactant.to_rarray(_host(t1)) + r2 = Reactant.to_rarray(_host(t2)) + compiled = @compile core(r1, r2) + return function (x, v) + out = compiled(Reactant.to_rarray(_host(x)), Reactant.to_rarray(_host(v))) + return _from_host(x, out) + end +end + +# Forward-mode JVP of `g` in direction `v`, i.e. J(g)·v. For g = gradlogp +# this is the HVP H·v. `Const(g)` so Enzyme doesn't treat captures as active. +function _jvp(g, x, v) + return only(Enzyme.autodiff(Enzyme.Forward, Enzyme.Const(g), Enzyme.Duplicated(x, v))) +end + +_rev_gradient(f, x) = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)[1] + +#= +Gradient slots. The HVP factories below get `logdensity` from `_resolve_hvp`, +which already has it. +=# +function ParallelMCMC._reactant_resolve_gradient( + logdensity, backend::AutoReactant, x_template::AbstractVector +) + _check_reactant_mode(backend) + core = Base.Fix1(_rev_gradient, logdensity) + return _compiled(core, x_template) +end + +function ParallelMCMC._reactant_resolve_gradient_batch( + logdensity_batch, backend::AutoReactant, X_template::AbstractMatrix +) + _check_reactant_mode(backend) + # `_BatchLogdensitySum` is the same sum-over-columns the DI-driven batched + # gradient uses (src/interface.jl), so both batched paths differentiate the + # same thing. + core = Base.Fix1(_rev_gradient, ParallelMCMC._BatchLogdensitySum(logdensity_batch)) + return _compiled(core, X_template) +end + +#= +HVP factories. `_resolve_hvp` / `_resolve_hvp_batch` (src/interface.jl) route to +one of two shapes, the same two every other backend gets. + + - Both `grad_logdensity` and the HVP source `AutoReactant`: the AD-derived + gradient case (see `_second_order` in `interface.jl`). Traced here as + forward-over-reverse from `logdensity`, a single XLA program. + - A hand-written `gradlogp` with an `AutoReactant` HVP source: routed by + `_hvp_strategy(::AutoReactant) = ReactantHVP()` in `DEER.jl` to + `_make_hvp_fn` below, a forward JVP over that callable. +=# +function DEER._make_hvp_fn_second_order( + logdensity, backend::AutoReactant, x_template::AbstractVector +) + _check_reactant_mode(backend) + inner = Base.Fix1(_rev_gradient, logdensity) + core(x, v) = _jvp(inner, x, v) + return _compiled(core, x_template, x_template) +end + +function DEER._make_hvp_fn( + ::DEER.ReactantHVP, gradlogp, backend::AutoReactant, x_template::AbstractVector +) + _check_reactant_mode(backend) + core(x, v) = _jvp(gradlogp, x, v) + return _compiled(core, x_template, x_template) +end + +function DEER._make_hvp_batch_fn_second_order( + logdensity_batch_sum, backend::AutoReactant, X_template::AbstractMatrix +) + _check_reactant_mode(backend) + inner = Base.Fix1(_rev_gradient, logdensity_batch_sum) + core(X, V) = _jvp(inner, X, V) + return _compiled(core, X_template, X_template) +end + +function DEER._make_hvp_batch_fn( + ::DEER.ReactantHVP, grad_batch, backend::AutoReactant, X_template::AbstractMatrix +) + # Column-independent batched gradient ⇒ forward JVP is the columnwise HVP. + _check_reactant_mode(backend) + core(X, V) = _jvp(grad_batch, X, V) + return _compiled(core, X_template, X_template) +end + +end # module diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index f0d677b..098300c 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -164,7 +164,7 @@ We bundle the closure with the prep so `prepare_gradient` and `gradient` see the same function instance (DI keys preparations on function identity). --------------------------------------------------------------------------- =# -import ..ParallelMCMC: pmcmc_dot, pmcmc_dotsum +import ..ParallelMCMC: pmcmc_dot, pmcmc_dotsum, _REACTANT_LOAD_HINT struct _HvpReverseClosure{F} grad::F @@ -177,17 +177,15 @@ 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, i.e., a hand-written -`gradlogp`, which neither of them differentiates twice: +Pick the AD-HVP fallback strategy from the user's backend. Both are one AD pass +over a hand-written `gradlogp`; an AD-derived gradient goes to +`_make_hvp_fn_second_order` instead. 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 @@ -203,32 +201,28 @@ abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end +# Need separate HVPStrategy for Reactant; DI does not support and so needs separate logic +struct ReactantHVP <: HVPStrategy end + _strategy_from(::DI.ForwardOverAnything) = ForwardOnGrad() _strategy_from(::DI.HVPMode) = ReverseOnGrad() _hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend)) +_hvp_strategy(::ADTypes.AutoReactant) = ReactantHVP() #= -Hook for backend-specific normalization of the user's `backend`, applied on every -AD-HVP path before the backend reaches DI. - -It supplies what the wrappers DEER differentiates need, and nothing else. Those -wrapper types are ours, so annotating them is ours to do: EnzymeExt specializes -this to fill `function_annotation=Enzyme.Const`, without which Enzyme throws -`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure`, which capture `gradlogp`. - -It deliberately does not choose a differentiation mode. Which direction a pass -runs is the user's call when they state one and DI's to resolve from the operator -when they don't; this package is not an AD package and has no business overriding -either. Picking one here also used to corrupt a `SecondOrder`, whose halves carry -directions of their own (see `_normalized_second_order`). - -Callers hand this a single pass, never a `SecondOrder`: the strategy paths below -run one AD pass over a hand-written `gradlogp`, and `_resolve_hvp` sends every -`SecondOrder` to `_make_hvp_fn_second_order` before they are reached. -`_normalized_second_order` is the one caller that starts from a pair, and it -selects the outer half itself. +Hook for backend-specific normalization, applied on every AD-HVP path before the +backend reaches DI. It fills in only what DEER's own wrapper types need; see +`ext/EnzymeExt.jl` for the one specialization that exists. + +It does not choose a differentiation mode. A mode the user set is a decision, an +unset one is DI's to resolve from the operator it runs, and substituting one here +used to rewrite the outer half of a `SecondOrder` out from under `hvp_mode` (see +`_normalized_second_order`). + +Only ever handed a single pass. `_resolve_hvp` sends every `SecondOrder` to +`_make_hvp_fn_second_order` before the strategy paths below are reached, and +`_normalized_second_order` picks the outer half itself. =# _normalized_backend(backend::AbstractADType) = backend @@ -300,23 +294,46 @@ function _make_hvp_batch_fn( return (X, V) -> _batch_hvp_via_grad_reverse_prepared(prep, X, V) end +# Fallbacks in case a user forgets `using Reactant` +function _make_hvp_fn( + ::ReactantHVP, gradlogp, backend::AbstractADType, x_template::AbstractVector +) + return error(_REACTANT_LOAD_HINT) +end + +function _make_hvp_batch_fn( + ::ReactantHVP, grad_batch, backend::AbstractADType, X_template::AbstractMatrix +) + return error(_REACTANT_LOAD_HINT) +end + +function _make_hvp_fn_second_order( + logdensity, backend::AbstractADType, x_template::AbstractVector +) + return error(_REACTANT_LOAD_HINT) +end + +function _make_hvp_batch_fn_second_order( + logdensity_batch_sum, backend::AbstractADType, X_template::AbstractMatrix +) + return error(_REACTANT_LOAD_HINT) +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. - -Both halves are passed to `DI.hvp` as the user composed them, so the direction -each one runs in is theirs and DI's, not ours. Normalization touches only the -outer half, and only to fill in annotations for the wrappers being -differentiated; because it never substitutes a mode, `DI.hvp_mode` of the pair is -the same before and after. The inner half is a plain first-order gradient over -the user's own `logdensity` and is passed straight through. - -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. +takes both passes over the log-density, so the gradient slot is never touched. +The alternative — pushing tangents through the prepared DI gradient — drops out +of that preparation the moment the outer pass hands it an unexpected tangent +type. + +Both halves reach `DI.hvp` as the user composed them. Normalization touches the +outer one, and only to fill in annotations, so `DI.hvp_mode` of the pair reads +the same before and after; the inner half is a first-order gradient over the +user's own `logdensity` and goes through untouched. + +The batched form differentiates `sum(logdensity_batch(X))`. Column independence +makes that Hessian block-diagonal, so its HVP along `V` is the columnwise HVP. --------------------------------------------------------------------------- =# function _normalized_second_order(backend::DI.SecondOrder) diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index e98eebb..126e64f 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -1,7 +1,7 @@ module ParallelMCMC using AbstractMCMC -using ADTypes: AbstractADType +using ADTypes: ADTypes, AbstractADType using CUDA using DifferentiationInterface: DifferentiationInterface using FlexiChains @@ -28,6 +28,10 @@ pmcmc_matmul(A::AbstractVecOrMat, B::AbstractVecOrMat) = A * B pmcmc_dot(a::AbstractVector, b::AbstractVector) = dot(a, b) pmcmc_dotsum(A::AbstractVecOrMat, B::AbstractVecOrMat) = sum(A .* B) +#= Lives here rather than in `DEER` because both DEER's `ReactantHVP` fallbacks +and `interface.jl`'s gradient hooks report it. =# +const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." + include("MALA/MALA.jl") include("DEER/DEERScan.jl") include("DEER/DEER.jl") diff --git a/src/interface.jl b/src/interface.jl index 69ede25..97aaec3 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -11,20 +11,28 @@ Defines model/sampler/state/transition types and implements Wraps a log-density function, its gradient, and optional Hessian-vector product helpers for use with ParallelMCMC samplers. -The derivative slots (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, -`hvp_batch`) take either a callable or an `ADTypes.AbstractADType`. Backends -are turned into prepared DifferentiationInterface callables when sampling -starts, and any AD failure surfaces there. So a model needs nothing beyond -the log-density: +Each derivative slot (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, +`hvp_batch`) takes a callable or an `ADTypes.AbstractADType`, so a model can be +built from the log-density alone: DensityModel(logp, AutoForwardDiff(), dim) -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 +Backends become prepared DifferentiationInterface callables when sampling starts, +and an AD failure surfaces there. `ADTypes.AutoReactant()` is the one backend DI +cannot drive; it is traced with Enzyme-MLIR and compiled to an XLA executable by +Reactant.jl instead, which needs `using Reactant` and brings requirements of its +own — see the GPU guide, `docs/src/15-gpu.md`, and `ext/ReactantExt.jl`'s module +docstring. + +A backend in `hvp` / `hvp_batch` over a hand-written gradient 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. +`AutoReactant` sits outside both: it cannot go inside a `SecondOrder`, and +`grad_logdensity` and `hvp`/`hvp_batch` must either both be `AutoReactant()` or +neither, since mixing it with a DifferentiationInterface backend across the two +passes of an HVP raises an `ArgumentError` at preparation time. - `logdensity(x::AbstractVector) -> Real` - `grad_logdensity` — callable `x -> AbstractVector`, or a backend to @@ -49,12 +57,11 @@ latter, and bypasses the gradient slot even when you wrote it by hand. information. 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. +update evaluates directly. `ParallelMALASampler` runs that update 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. A +`logdensity_batch` on its own is fine too: it scores whole trajectories at once +and leaves the batched update off. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel logdensity::F @@ -95,15 +102,10 @@ function DensityModel( "grad_logdensity must be a callable or an ADTypes.AbstractADType backend" ), ) - #= 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. =# + # `_prepare_model` decides from the gradient slot whether the batched path runs. _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", + "$name requires logdensity_batch, which the batched DEER update evaluates" ), ) if logdensity_batch === nothing @@ -123,12 +125,13 @@ function DensityModel( end """ -A [`DensityModel`](@ref) with its backend slots resolved to prepared -DifferentiationInterface callables. `_prepare_model` builds these, and the -sampler internals take them rather than a `DensityModel`, so no slot of one -of these ever holds an `AbstractADType`. Which slots are filled depends on -the sampler: the sequential samplers only need `grad_logdensity` and get -`nothing` for the DEER-only slots, DEER fills the rest. +A [`DensityModel`](@ref) with its backend slots resolved to prepared callables — +DifferentiationInterface ones, or a compiled Reactant/XLA executable for +`AutoReactant()`. `_prepare_model` builds these, and the sampler internals take +them rather than a `DensityModel`, so no slot here ever holds an +`AbstractADType`. Which slots are filled depends on the sampler: the sequential +samplers only need `grad_logdensity` and get `nothing` for the DEER-only slots, +DEER fills the rest. `source` is the `DensityModel` this was prepared from. A sampler state carries a prepped model so the preparation is reused across steps, and `initial_state` @@ -147,19 +150,16 @@ struct PreppedDensityModel{F,G,H,FB,GB,HB,PN,SM<:DensityModel} source::SM end -#= -Whether a prepped model carried in a state was built from the model a `step` +#= Whether a prepped model carried in a state was built from the model a `step` was handed. Identity, not equality: an equal-but-distinct `DensityModel` just -costs one re-preparation, whereas treating a different model as a match would -silently sample the wrong target. -=# +costs one re-preparation, whereas a false match would 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` rather than failing. In a -normal run the prepared branch is the one that fires. +anything else falls back to an unprepared `DI.gradient` rather than failing. =# struct _ADGradient{F,B<:AbstractADType,P,TX} logdensity::F @@ -211,6 +211,25 @@ function _resolve_gradient_batch( ) end +# `AutoReactant` gradients bypass DI +function _resolve_gradient( + logdensity, backend::ADTypes.AutoReactant, x_template::AbstractVector +) + return _reactant_resolve_gradient(logdensity, backend, x_template) +end +function _reactant_resolve_gradient(logdensity, backend, x_template) + return error(_REACTANT_LOAD_HINT) +end + +function _resolve_gradient_batch( + logdensity_batch, backend::ADTypes.AutoReactant, X_template::AbstractMatrix +) + return _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) +end +function _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) + return error(_REACTANT_LOAD_HINT) +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 @@ -226,7 +245,7 @@ function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) 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 + logdensity, _second_order(hvp_backend, grad_backend), x_template ) else return DEER._make_hvp_fn( @@ -239,6 +258,7 @@ end function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) + _check_reactant_pair(grad_batch_backend, hvp_backend) if hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template @@ -246,7 +266,7 @@ function _resolve_hvp_batch( elseif grad_batch_backend !== nothing return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), - DI.SecondOrder(hvp_backend, grad_batch_backend), + _second_order(hvp_backend, grad_batch_backend), X_template, ) else @@ -256,6 +276,90 @@ function _resolve_hvp_batch( end end +#= The HVP backend composed with the backend that produced the gradient under it, +via `DI.SecondOrder(hvp_backend, grad_backend)` and `DI.hvp`. Two `AutoReactant`s +are not a pair DI could run, so they collapse to `hvp_backend`: the backend +`ReactantExt` traces forward-over-reverse from `logdensity` in +`_make_hvp_fn_second_order`. Returning it, not a fresh `AutoReactant()`, keeps a +non-default `mode` reachable by `_check_reactant_mode`. A mixed pair is already +out by the time this runs, via `_check_reactant_pair`. =# +_second_order(hvp_backend, grad_backend) = DI.SecondOrder(hvp_backend, grad_backend) +function _second_order(hvp_backend::ADTypes.AutoReactant, ::ADTypes.AutoReactant) + return hvp_backend +end + +# Check Reactant isn't passed with DI +_check_reactant_pair(grad_backend, hvp_backend) = nothing +_check_reactant_pair(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = nothing +_check_reactant_pair(::Nothing, ::ADTypes.AutoReactant) = nothing + +function _check_reactant_pair(grad_backend::ADTypes.AutoReactant, hvp_backend) + return throw( + ArgumentError( + "an AutoReactant gradient needs an AutoReactant Hessian-vector product: " * + "got hvp backend $(hvp_backend). Reactant compiles the gradient to an XLA " * + "executable, which DifferentiationInterface cannot differentiate. Set the " * + "model's `hvp` (or the sampler's `backend`) to AutoReactant() as well.", + ), + ) +end + +function _check_reactant_pair(grad_backend, hvp_backend::ADTypes.AutoReactant) + return throw( + ArgumentError( + "an AutoReactant Hessian-vector product needs an AutoReactant or " * + "hand-written gradient: got gradient backend $(grad_backend). Reactant " * + "traces the HVP from the log-density (or from your gradient) and cannot " * + "trace a DifferentiationInterface-prepared gradient. Set " * + "`grad_logdensity` to AutoReactant() or supply a callable.", + ), + ) +end + +#= Would otherwise be ambiguous between the two methods above, and wants its own +message anyway: "an AutoReactant Hessian-vector product" does not describe a +`SecondOrder`. =# +function _check_reactant_pair(::ADTypes.AutoReactant, hvp_backend::DI.SecondOrder) + return throw( + ArgumentError( + "an AutoReactant gradient needs a bare AutoReactant Hessian-vector " * + "product: got hvp backend $(hvp_backend). An AutoReactant gradient " * + "compiles to an XLA executable, which DifferentiationInterface's " * + "SecondOrder cannot drive. Set the model's `hvp` (or the sampler's " * + "`backend`) to AutoReactant(), unwrapped.", + ), + ) +end + +#= `SecondOrder(AutoReactant(), AutoReactant())` is a natural thing to try, given +the pairing table in `10-getting-started.md`, and would otherwise land in +`DI.prepare_hvp` several frames deep with no Reactant support. Checked whatever +the gradient slot holds, since a `SecondOrder` bypasses it anyway (see +`_resolve_hvp`). =# +function _check_reactant_pair(grad_backend, hvp_backend::DI.SecondOrder) + _second_order_has_reactant(hvp_backend) || return nothing + return throw( + ArgumentError( + "AutoReactant cannot go inside a DifferentiationInterface SecondOrder: " * + "got hvp backend $(hvp_backend). DifferentiationInterface has no Reactant " * + "support at all. Set `hvp` (or the sampler's `backend`) to a bare " * + "AutoReactant() instead.", + ), + ) +end + +function _second_order_has_reactant(so::DI.SecondOrder) + return DI.outer(so) isa ADTypes.AutoReactant || DI.inner(so) isa ADTypes.AutoReactant +end + +#= A `LogDensityProblemGradient` (defined below) is a callable, so it clears the +`grad_backend === nothing` test that otherwise means "hand-written gradient" — +but it dispatches into DynamicPPL/LogDensityProblems and is not +Reactant-traceable. `_check_reactant_pair` only sees the backend, which is +`nothing` for both, so this checks `grad`'s type instead. The specialization has +to wait for `LogDensityProblemGradient` to exist and sits further down. =# +_check_reactant_hvp_source(grad, hvp_backend) = nothing + """ _prepare_model(model, x_template) -> PreppedDensityModel _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel @@ -281,10 +385,8 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector) else model.grad_logdensity end - #= The derivative slots DEER alone reads are dropped rather than passed - along: a sequential sampler never looks at them, and carrying an - unresolved backend would break the invariant that no slot here holds an - `AbstractADType`. =# + # Dropped, not passed along: a sequential sampler never reads them, and an + # unresolved backend there would break `PreppedDensityModel`'s invariant. return PreppedDensityModel( model.logdensity, grad, @@ -301,31 +403,42 @@ end function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, backend) 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 === nothing || model.hvp isa AbstractADType - hvp_backend = model.hvp === nothing ? backend : model.hvp - hvp_backend === nothing && throw( + #= Settle the HVP backend and check its pairing with `grad_backend` before + resolving the gradient. Otherwise a mismatched `AutoReactant` pair, or the + LogDensityProblems-gradient case, surfaces only once `_resolve_gradient` has + paid for an XLA compile: 18+ seconds to report a config error. =# + needs_hvp = model.hvp === nothing || model.hvp isa AbstractADType + hvp_backend = if needs_hvp + hb = model.hvp === nothing ? backend : model.hvp + hb === nothing && throw( ArgumentError( "ParallelMALASampler needs a Hessian-vector product: supply `hvp` " * "on the DensityModel (callable or AD backend), or pass `backend=` " * "to ParallelMALASampler", ), ) + _check_reactant_pair(grad_backend, hb) + _check_reactant_hvp_source(model.grad_logdensity, hb) + hb + else + nothing + end + + grad = if grad_backend !== nothing + _resolve_gradient(model.logdensity, grad_backend, x_template) + else + model.grad_logdensity + end + + hvp = if needs_hvp _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. =# + # Leaves the batched path off rather than raising: `_trajectory_logps` uses + # `logdensity_batch` either way. grad_batch = model.grad_logdensity_batch if grad_batch === nothing && model.logdensity_batch !== nothing grad_batch = grad_backend @@ -340,26 +453,35 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, X_template = similar(x_template, length(x_template), T) X_template .= x_template - if grad_batch_backend !== nothing - grad_batch = _resolve_gradient_batch( - model.logdensity_batch, grad_batch_backend, X_template - ) - end - - if hvp_batch === nothing || hvp_batch isa AbstractADType + # Same reasoning as the unbatched case above: check before compiling. + needs_hvp_batch = hvp_batch === nothing || hvp_batch isa AbstractADType + hvp_batch_backend = if needs_hvp_batch # The model's own HVP backend if it has one, else the sampler's. - hvp_batch_backend = if hvp_batch === nothing + hbb = if hvp_batch === nothing model.hvp isa AbstractADType ? model.hvp : backend else hvp_batch end - hvp_batch_backend === nothing && throw( + hbb === 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", ), ) + _check_reactant_pair(grad_batch_backend, hbb) + hbb + else + nothing + end + + if grad_batch_backend !== nothing + grad_batch = _resolve_gradient_batch( + model.logdensity_batch, grad_batch_backend, X_template + ) + end + + if needs_hvp_batch hvp_batch = _resolve_hvp_batch( model.logdensity_batch, grad_batch, @@ -405,6 +527,20 @@ struct LogDensityProblemGradient{L} ld::L end +# The `_check_reactant_hvp_source` specialization promised further up. +function _check_reactant_hvp_source( + ::LogDensityProblemGradient, hvp_backend::ADTypes.AutoReactant +) + return throw( + ArgumentError( + "an AutoReactant Hessian-vector product needs an AutoReactant or " * + "hand-written gradient: got a LogDensityProblems-derived gradient. Reactant " * + "cannot trace DynamicPPL/LogDensityProblems machinery. Set `grad_logdensity` " * + "to AutoReactant(), or supply a Reactant-traceable callable.", + ), + ) +end + """ MALASampler(epsilon; cholM=nothing) @@ -581,11 +717,15 @@ 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 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. +`backend` supplies Hessian-vector products when the `DensityModel` brings no +`hvp` / `hvp_batch` of its own, and does nothing else. It never supplies a +gradient, so it cannot decide which update path runs, nor put AD on a function +the model had no backend for. Leave it out for a model that carries its own HVPs. + +`backend = ADTypes.AutoReactant()` constrains the model too, since +`AutoReactant` does not pair with a DifferentiationInterface backend across the +two passes of an HVP: `grad_logdensity` must then be `AutoReactant()` as well, or +a hand-written callable. Mixing raises an `ArgumentError` at preparation time. """ struct ParallelMALASampler{FP<:AbstractFloat,CM,AD} <: AbstractMCMC.AbstractSampler epsilon::FP @@ -934,9 +1074,9 @@ function _construct_flexichain( param_names::Any, model::DensityModel, ) where {TKey} - #= Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., - `param_names=(:x, :y, :z=>(2,))` without faffing with `Parameter` themselves. Also - 'upgrade' symbol parameter names to VarNames if the user requested a VNChain. =# + # Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., + # `param_names=(:x, :y, :z=>(2,))` without faffing with `Parameter` themselves. Also + # 'upgrade' symbol parameter names to VarNames if the user requested a VNChain. to_parameter(vn::VarName) = FlexiChains.Parameter(vn) to_parameter(s::Symbol) = FlexiChains.Parameter(TKey <: VarName ? VarName{s}() : s) diff --git a/test/Project.toml b/test/Project.toml index 7dbcce5..a2370e2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -22,8 +22,9 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +[compat] +JuliaFormatter = "2.12" + [extras] CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" - -[compat] -JuliaFormatter = "2.12" \ No newline at end of file +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" diff --git a/test/runtests.jl b/test/runtests.jl index 45cecaf..4d5b1d8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,13 @@ using ParallelMCMC using Test +#= `Reactant_jll` ships a prebuilt XLA and is a large download, so `Reactant` sits +in test/Project.toml's `[extras]` and `test-Reactant-HVP.jl` runs only when this +is set. Add Reactant to the test environment as well; the file's own +`try ... using Reactant ... catch` skips its testsets if it still won't load. =# +const _RUN_REACTANT_TESTS = + lowercase(get(ENV, "PARALLELMCMC_TEST_REACTANT", "false")) in ("1", "true", "yes") + @testset verbose=true "ParallelMCMC" begin #= Don't add your tests to runtests.jl. Instead, create files named @@ -14,6 +21,10 @@ using Test if isnothing(match(r"^test-.*\.jl$", file)) continue end + if file == "test-Reactant-HVP.jl" && !_RUN_REACTANT_TESTS + @info "Skipping $file (set PARALLELMCMC_TEST_REACTANT=true to opt in)" + continue + end title = titlecase(replace(splitext(file[6:end])[1], "-" => " ")) @testset verbose=true "$title" begin include(joinpath(root, file)) # robust if walkdir recurses diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index e56997e..44305c1 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -24,6 +24,14 @@ const DI_STRAT = ParallelMCMC.DEER.DI DEER_STRAT.ReverseOnGrad end + @testset "AutoReactant short-circuits hvp_mode" begin + #= DI cannot drive Reactant, so `AutoReactant` never reaches `DI.hvp_mode` + and the strategy is picked by dispatch instead. The method lives in DEER + rather than in ReactantExt, so this holds with Reactant unloaded. =# + @test DEER_STRAT._hvp_strategy(AutoReactant()) isa DEER_STRAT.ReactantHVP + @test @inferred(DEER_STRAT._hvp_strategy(AutoReactant())) isa DEER_STRAT.ReactantHVP + end + @testset "SecondOrder follows hvp_mode's composition" begin so_fwd_outer = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()) @test DEER_STRAT._hvp_strategy(so_fwd_outer) isa DEER_STRAT.ForwardOnGrad @@ -38,8 +46,7 @@ const DI_STRAT = ParallelMCMC.DEER.DI @testset "normalization supplies Const but never a mode" begin #= The wrappers DEER differentiates are its own types, so annotating them - `Const` is its business. The mode is not: one the user set is a decision, - and an unset one is DI's to resolve from the operator it runs. =# + `Const` is its business. The mode is not. =# bare = DEER_STRAT._normalized_backend(AutoEnzyme()) @test bare isa AutoEnzyme{<:Any,Enzyme.Const} @test bare.mode === nothing @@ -62,9 +69,9 @@ const DI_STRAT = ParallelMCMC.DEER.DI @testset "normalizing a SecondOrder keeps the composition DI resolved" begin #= Regression for #62. Normalization used to route the outer half through a forward-only hook, which pinned `Enzyme.Forward` onto it. For a pair - `hvp_mode` resolves to reverse -- `SecondOrder(AutoEnzyme(), + `hvp_mode` resolves to reverse — `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is reverse-over-forward, its inner half being - forward-only -- that silently made it forward-over-forward. =# + forward-only — that made it forward-over-forward. =# for so in ( DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()), DI_STRAT.SecondOrder(AutoEnzyme(), AutoZygote()), diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl new file mode 100644 index 0000000..e85a3e2 --- /dev/null +++ b/test/test-Reactant-HVP.jl @@ -0,0 +1,393 @@ +using Test +using Random +using LinearAlgebra +using Statistics +using FlexiChains + +using ParallelMCMC +using ADTypes +# Stands in for "some backend that is not Reactant" in the pairing tests below. +using ForwardDiff: ForwardDiff +using LogDensityProblems: LogDensityProblems + +const DI_R = ParallelMCMC.DEER.DI + +#= +Reactant-compiled derivative paths (`AutoReactant`), see ext/ReactantExt.jl. + +The quartic target is needed to test second order structure: logp = -0.25‖x‖⁴ has +H = -(‖x‖² I + 2 x xᵀ), so an HVP that drops the second-order term is caught, +where a Gaussian's constant H would hide it. Every derivative-accuracy testset +below uses it; do not swap in a Gaussian. +=# +logp_r(x) = -0.25 * sum(abs2, x)^2 +gradlogp_r(x) = -sum(abs2, x) .* x +hvp_r(x, v) = -(sum(abs2, x) .* v .+ 2 .* dot(x, v) .* x) +logp_batch_r(X) = vec(-0.25 .* sum(abs2, X; dims=1) .^ 2) +gradlogp_batch_r(X) = -X .* sum(abs2, X; dims=1) +logp_r32(x) = -0.25f0 * sum(abs2, x)^2 + +#= Standard Gaussian, constant Hessian (H = -I). Only for the end-to-end sampling +tests, which ask whether the sampler converged to the right posterior; whether +the HVP is second-order-correct is the quartic target's job. =# +logp_gauss(x) = -0.5 * sum(abs2, x) +gradlogp_gauss(x) = -x +hvp_gauss(x, v) = -v + +const D_R = 4 +const CT_R = FlexiChains.FlexiChain{Symbol} + +#= The pairing rule lives in `_resolve_hvp` / `_prepare_model`, not in the +extension, so its dispatch table is checked whether or not Reactant loads. +None of these calls resolve a gradient or compile anything. =# +@testset "Reactant does not pair with a DI backend" begin + # Both slots Reactant, or a hand-written gradient (`nothing`), are accepted. + @test ParallelMCMC._check_reactant_pair(AutoReactant(), AutoReactant()) === nothing + @test ParallelMCMC._check_reactant_pair(nothing, AutoReactant()) === nothing + @test ParallelMCMC._check_reactant_pair(AutoForwardDiff(), AutoForwardDiff()) === + nothing + + # One of each is refused in both directions. + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoReactant(), AutoForwardDiff() + ) + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoForwardDiff(), AutoReactant() + ) +end + +#= `SecondOrder(AutoReactant(), AutoReactant())`, or `AutoReactant()` paired with +a `SecondOrder` at all, is a natural thing to try given the pairing table in +10-getting-started.md, and would otherwise land in `DI.prepare_hvp` several +frames deep with no Reactant support. A dispatch-table property, so this runs +whether or not Reactant is loaded. =# +@testset "AutoReactant cannot appear inside a SecondOrder" begin + # Sanity: an ordinary SecondOrder is unaffected. + @test ParallelMCMC._check_reactant_pair( + nothing, DI_R.SecondOrder(AutoForwardDiff(), AutoForwardDiff()) + ) === nothing + + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + nothing, DI_R.SecondOrder(AutoReactant(), AutoReactant()) + ) + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoForwardDiff(), DI_R.SecondOrder(AutoReactant(), AutoForwardDiff()) + ) + # (AutoReactant, SecondOrder) exercises the disambiguating method directly. + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoReactant(), DI_R.SecondOrder(AutoReactant(), AutoReactant()) + ) +end + +#= A `LogDensityProblemGradient` is a callable, so it clears the +`grad_backend === nothing` test that otherwise means "hand-written gradient", +but it is not Reactant-traceable. Checked on the type first, then through a real +LogDensityProblems model. Neither needs Reactant loaded: the check fires before +any gradient is resolved or anything compiled. =# +@testset "a LogDensityProblems gradient cannot pair with an AutoReactant hvp" begin + @test ParallelMCMC._check_reactant_hvp_source( + ParallelMCMC.LogDensityProblemGradient(nothing), AutoForwardDiff() + ) === nothing + @test_throws ArgumentError ParallelMCMC._check_reactant_hvp_source( + ParallelMCMC.LogDensityProblemGradient(nothing), AutoReactant() + ) + + struct _FakeLD end + LogDensityProblems.capabilities(::_FakeLD) = LogDensityProblems.LogDensityOrder{1}() + LogDensityProblems.dimension(::_FakeLD) = D_R + function LogDensityProblems.logdensity_and_gradient(::_FakeLD, x) + return logp_r(x), gradlogp_r(x) + end + + model = DensityModel(_FakeLD(); hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(model, zeros(D_R), 8, nothing) +end + +#= Outside the `reactant_ok` guard below. `_check_reactant_pair` now runs above +gradient resolution in `_prepare_model`, so a mismatched pair is caught before +anything is resolved or compiled and no Reactant install is needed. This used to +pay a full XLA compile per `@test_throws`, ~18s, to demonstrate a config error. =# +@testset "mixed pairs are refused at preparation" begin + x = zeros(D_R) + + reactant_grad = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoForwardDiff()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_grad, x, 8, nothing) + + reactant_hvp = DensityModel(logp_r, AutoForwardDiff(), D_R; hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_hvp, x, 8, nothing) +end + +reactant_ok = try + using Reactant: Reactant + using Enzyme: Enzyme + true +catch err + @warn "Reactant not available — skipping Reactant HVP tests" err + false +end + +#= The `_REACTANT_LOAD_HINT` fallbacks (src/DEER/DEER.jl, and +`_reactant_resolve_gradient` / `_reactant_resolve_gradient_batch` in +src/interface.jl) can only be asserted in a session without Reactant: once +`ReactantExt` is loaded its more-specific methods shadow all of them and the +error can never fire. Hence the one testset here guarded on `!reactant_ok`. =# +if !reactant_ok + @testset "clear load-hint error without Reactant loaded" begin + #= Gradient slot: `_reactant_resolve_gradient`'s fallback. `backend` (the + 4th arg) is `AutoReactant()` too, as `_check_reactant_pair` requires, so + `_prepare_model` gets as far as gradient resolution rather than failing + first on the model having no HVP source — a correct failure, but not the + one being pinned here. =# + model_grad = DensityModel(logp_r, AutoReactant(), D_R) + @test_throws "AutoReactant requires Reactant.jl" ParallelMCMC._prepare_model( + model_grad, zeros(D_R), 8, AutoReactant() + ) + + # HVP slot over a hand-written gradient: `DEER._make_hvp_fn`'s + # `ReactantHVP` fallback, reached via `_hvp_strategy(::AutoReactant)`. + model_hvp = DensityModel(logp_r, gradlogp_r, D_R; hvp=AutoReactant()) + @test_throws "AutoReactant requires Reactant.jl" ParallelMCMC._prepare_model( + model_hvp, zeros(D_R), 8, nothing + ) + end +end + +if reactant_ok + @testset "extension is loaded" begin + @test Base.get_extension(ParallelMCMC, :ReactantExt) !== nothing + end + + #= The extension does not honour `AutoReactant.mode` (the wrapped + `AutoEnzyme`): gradients always trace reverse, HVPs forward-over-that. A + non-default mode is rejected rather than ignored. =# + @testset "a non-default AutoReactant mode is rejected" begin + bad = AutoReactant(; mode=AutoEnzyme(; mode=Enzyme.Forward)) + model_grad = DensityModel(logp_r, bad, D_R) + @test_throws ArgumentError ParallelMCMC._prepare_model( + model_grad, zeros(D_R), 8, nothing + ) + + model_hvp = DensityModel(logp_r, gradlogp_r, D_R; hvp=bad) + @test_throws ArgumentError ParallelMCMC._prepare_model( + model_hvp, zeros(D_R), 8, nothing + ) + + # The default still works. + @test DensityModel(logp_r, AutoReactant(), D_R) isa DensityModel + end + + @testset "HVP matches analytic" begin + rng = MersenneTwister(71) + x = randn(rng, D_R) + v = randn(rng, D_R) + + @testset "forward over user gradient (sampler backend)" begin + model = DensityModel(logp_r, gradlogp_r, D_R) + m_p = ParallelMCMC._prepare_model(model, x, 8, AutoReactant()) + @test m_p.hvp(x, v) ≈ hvp_r(x, v) + end + + @testset "forward-over-reverse from logp alone (both slots AutoReactant)" begin + model = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x, 8, nothing) + @test m_p.grad_logdensity(x) ≈ gradlogp_r(x) + @test m_p.hvp(x, v) ≈ hvp_r(x, v) + end + + @testset "batched slots" begin + T = 8 + X = randn(rng, D_R, T) + V = randn(rng, D_R, T) + Hv_cols = reduce(hcat, [hvp_r(X[:, t], V[:, t]) for t in 1:T]) + + model = DensityModel( + logp_r, + AutoReactant(), + D_R; + logdensity_batch=logp_batch_r, + grad_logdensity_batch=AutoReactant(), + hvp=AutoReactant(), + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, X[:, 1], T, nothing) + @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) + @test m_p.hvp_batch(X, V) ≈ Hv_cols + end + + #= A hand-written batched gradient with `hvp_batch=AutoReactant()`, the + batched analogue of the "forward over user gradient" case above. Routes + through + `DEER._make_hvp_batch_fn(::ReactantHVP, grad_batch, ::AutoReactant, ...)`, + which nothing else here reaches. =# + @testset "forward over user batched gradient (hvp_batch=AutoReactant())" begin + T = 8 + X = randn(rng, D_R, T) + V = randn(rng, D_R, T) + Hv_cols = reduce(hcat, [hvp_r(X[:, t], V[:, t]) for t in 1:T]) + + model = DensityModel( + logp_r, + gradlogp_r, + D_R; + hvp=hvp_r, + logdensity_batch=logp_batch_r, + grad_logdensity_batch=gradlogp_batch_r, + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, X[:, 1], T, nothing) + @test m_p.hvp_batch(X, V) ≈ Hv_cols + end + + #= Edge shapes. Both slots AutoReactant, as in the + "forward-over-reverse from logp alone" case above, but at the smallest + sizes DEER ever prepares. =# + @testset "edge shapes" begin + @testset "D=1" begin + x1 = randn(rng, 1) + v1 = randn(rng, 1) + model = DensityModel(logp_r, AutoReactant(), 1; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x1, 8, nothing) + @test m_p.grad_logdensity(x1) ≈ gradlogp_r(x1) + @test m_p.hvp(x1, v1) ≈ hvp_r(x1, v1) + end + + @testset "T=1" begin + T = 1 + X = reshape(x, D_R, T) + V = reshape(v, D_R, T) + model = DensityModel( + logp_r, + AutoReactant(), + D_R; + logdensity_batch=logp_batch_r, + grad_logdensity_batch=AutoReactant(), + hvp=AutoReactant(), + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, x, T, nothing) + @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) + @test m_p.hvp_batch(X, V) ≈ reshape(hvp_r(x, v), D_R, T) + end + end + + #= Float32 on the CPU path. Previously only exercised inside the + CuArray-only block below, so it never ran without a functional CUDA. =# + @testset "Float32 on CPU" begin + x32 = Float32.(x) + v32 = Float32.(v) + model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x32, 8, nothing) + + g = m_p.grad_logdensity(x32) + @test eltype(g) === Float32 + @test g ≈ gradlogp_r(x32) + + Hv = m_p.hvp(x32, v32) + @test eltype(Hv) === Float32 + @test Hv ≈ hvp_r(x32, v32) + end + end + + #= Pins the "captured data is frozen at preparation time" limitation from + ext/ReactantExt.jl's module docstring and docs/src/15-gpu.md: `@compile` + bakes a captured plain `Array` in as a compile-time constant, so mutating it + afterwards has NO effect on later calls, with no error and no warning. Here + so that a future Reactant which does detect this gets noticed. Model code + still should not close over mutable data. + =# + @testset "documented caveat: captured data is frozen at preparation time" begin + data = [1.0, 1.0, 1.0, 1.0] + f(x) = -0.5 * sum(abs2, x .- data) # ∇f(x) = data - x + model = DensityModel(f, AutoReactant(), D_R) + # Two-argument `_prepare_model` only resolves `grad_logdensity`, which + # is all this test needs (no HVP involved). + m_p = ParallelMCMC._prepare_model(model, zeros(D_R)) + + x0 = zeros(D_R) + g_before = m_p.grad_logdensity(x0) + @test g_before ≈ [1.0, 1.0, 1.0, 1.0] + + data .= 5.0 # mutate the captured array *after* preparation + + g_after = m_p.grad_logdensity(x0) + #= Frozen at the pre-mutation value: NOT [5, 5, 5, 5], which is what a + correct re-evaluation against the mutated `data` would give. =# + @test g_after ≈ [1.0, 1.0, 1.0, 1.0] + end + + #= `size(chain) == (N,1)` and `all(isfinite, ...)` pass even for a badly + wrong HVP: DEER's Newton iteration just fails to converge and `DEER.solve` + returns the non-converged trajectory, no NaNs involved. Assert the posterior + mean against a target with a known mean instead (the convention in + test-GPU-AD-HVP.jl), and cross-check against the same model driven by the + analytic HVP on the same noise tape. + =# + @testset "end-to-end sampling with AutoReactant" begin + @testset "posterior mean recovery (standard Gaussian, mean 0)" begin + model = DensityModel(logp_gauss, AutoReactant(), D_R) + s = ParallelMALASampler(0.3; T=16, backend=AutoReactant()) + n_samples, n_burn = 2000, 500 + chain = sample( + MersenneTwister(72), model, s, n_samples; chain_type=CT_R, progress=false + ) + @test size(chain) == (n_samples, 1) + + xs = chain[:x] + @test all(x -> all(isfinite, x), xs) + post_mean = vec(mean(reduce(hcat, xs[(n_burn + 1):end]); dims=2)) + @test maximum(abs, post_mean) < 0.3 + end + + @testset "matches analytic-HVP DEER on the same noise tape" begin + s = ParallelMALASampler(0.1; T=16, backend=AutoReactant()) + model_r = DensityModel(logp_gauss, AutoReactant(), D_R) + model_an = DensityModel(logp_gauss, gradlogp_gauss, D_R; hvp=hvp_gauss) + + c_r = sample( + MersenneTwister(99), model_r, s, 64; chain_type=CT_R, progress=false + ) + c_an = sample( + MersenneTwister(99), model_an, s, 64; chain_type=CT_R, progress=false + ) + @test c_r[:x] ≈ c_an[:x] + end + end + + reactant_gpu_ok = try + using CUDA: CUDA + CUDA.functional() && (CUDA.CuArray([1.0f0]); true) + catch + false + end + + if !reactant_gpu_ok + @info "Reactant HVP test: CUDA not functional — skipping CuArray boundary" + else + #= + The compiled executable lives in Reactant's own (XLA) device memory, so + what this checks is the CuArray <-> Reactant marshalling boundary: + CuArray in, CuArray out, values matching the analytic HVP. It does NOT + establish that the HVP executes on the GPU, which depends on Reactant's + default XLA client (see docs/src/15-gpu.md) and is neither controlled + nor asserted here. + =# + @testset "CuArray boundary" begin + rng = MersenneTwister(73) + x_h = randn(rng, Float32, D_R) + v_h = randn(rng, Float32, D_R) + x_d = CUDA.CuArray(x_h) + v_d = CUDA.CuArray(v_h) + + model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x_d, 8, nothing) + + g = m_p.grad_logdensity(x_d) + @test g isa CUDA.CuArray + @test Array(g) ≈ gradlogp_r(x_h) + + Hv = m_p.hvp(x_d, v_d) + @test Hv isa CUDA.CuArray + @test Array(Hv) ≈ hvp_r(x_h, v_h) + end + end +end