From 16a77af95a076b3c9e829e50e29a10514705de96 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:54:55 -0400 Subject: [PATCH 1/2] Drop Turing/DynamicPPL support DynamicPPL models only run on the CPU, which is the case ParallelMCMC is least suited for, and keeping the extension working is more maintenance than it is worth right now. Releases up to 0.2 keep it. Removes DynamicPPLExt and its tests, strips the Turing model from the logistic-regression DEER tests, and adds tests for LogDensityProblemsExt, which were previously only exercised through DynamicPPL. Closes #74 --- CHANGELOG.md | 9 +- Project.toml | 3 - README.md | 4 +- docs/src/10-getting-started.md | 54 +-- docs/src/15-gpu.md | 16 +- docs/src/95-reference.md | 7 +- docs/src/index.md | 23 -- ext/DynamicPPLExt.jl | 167 -------- ext/LogDensityProblemsExt.jl | 30 +- src/interface.jl | 3 +- test/Project.toml | 2 - ...ring-Logistic.jl => test-DEER-Logistic.jl} | 96 +---- test/test-GPU-AD-HVP.jl | 2 +- test/test-LogDensityProblems.jl | 81 ++++ test/test-Turing-Integration.jl | 384 ------------------ 15 files changed, 119 insertions(+), 762 deletions(-) delete mode 100644 ext/DynamicPPLExt.jl rename test/{test-DEER-Turing-Logistic.jl => test-DEER-Logistic.jl} (70%) create mode 100644 test/test-LogDensityProblems.jl delete mode 100644 test/test-Turing-Integration.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 2afa38b..a98d062 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,8 +22,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `ADTypes.AutoReactant()` can supply derivatives and sampler HVPs through Reactant.jl. It requires `using Reactant`, a traceable log-density, and `AutoReactant()` on both AD passes of an HVP (#37, #52). -- The `DensityModel` constructors in `DynamicPPLExt` and `LogDensityProblemsExt` - now accept `hvp` and the batched derivative slots. +- The `DensityModel` constructor in `LogDensityProblemsExt` now accepts `hvp` + and the batched derivative slots. ### Fixed @@ -54,7 +54,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Removed -- `DynamicPPLExt` no longer requires `ForwardDiff` as a triggering library to load. +- Turing/DynamicPPL support (#74). `DynamicPPLExt` and the + `DensityModel(::DynamicPPL.Model)` constructor are gone, since DynamicPPL + models cannot run on the GPU. Use ParallelMCMC 0.2 for Turing models, or wrap + the log-density in a `LogDensityProblems` object or a plain `DensityModel`. ## [0.2.0] - 2026-06-29 diff --git a/Project.toml b/Project.toml index cbc6af3..ac2df19 100644 --- a/Project.toml +++ b/Project.toml @@ -15,7 +15,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [weakdeps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" @@ -25,7 +24,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] CUDAExt = "CUDA" -DynamicPPLExt = ["DynamicPPL", "FlexiChains", "LogDensityProblems"] EnzymeExt = "Enzyme" LogDensityProblemsExt = "LogDensityProblems" ReactantExt = ["Reactant", "Enzyme"] @@ -36,7 +34,6 @@ AbstractMCMC = "5.10.0" CUDA = "5.11.0" CUDA_Runtime_jll = "0.21" DifferentiationInterface = "0.7.13" -DynamicPPL = "0.41.6, 0.42" Enzyme = "0.13.146" FlexiChains = "0.6.6" ForwardDiff = "1" diff --git a/README.md b/README.md index 6398069..b443e61 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ The approach and its scaling tricks (stochastic Hutchinson Jacobian estimators, | [`MALASampler`](src/interface.jl) | Baseline — sequential MALA with a fixed step size | | [`AdaptiveMALASampler`](src/interface.jl) | Baseline — sequential MALA with dual-averaging step-size adaptation | -All samplers implement the [AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and return [`FlexiChains`](https://pysm.dev/FlexiChains.jl) objects, so they slot into existing Turing.jl / AbstractMCMC workflows. +All samplers implement the [AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and return [`FlexiChains`](https://pysm.dev/FlexiChains.jl) objects, so they slot into existing AbstractMCMC workflows. ### Quick start @@ -66,7 +66,7 @@ sampler = ParallelMALASampler(0.1; T=64, jacobian=:stoch_diag, chain = sample(model, sampler, 500; chain_type=VNChain) ``` -See the [Getting Started guide](docs/src/10-getting-started.md) for worked examples, GPU usage, Turing.jl integration, and step-size tuning. +See the [Getting Started guide](docs/src/10-getting-started.md) for worked examples, GPU usage, and step-size tuning. ## How to Cite diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index 64be447..56269d2 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -47,7 +47,7 @@ Hand-written and AD-derived slots can be mixed. | backend | backend | `SecondOrder(hvp, grad_logdensity)` on `logdensity` | | either | `SecondOrder(...)` | that pair on `logdensity`, gradient slot unused | -An explicit `SecondOrder` bypasses the gradient slot. Use it for an AD-derived HVP on Turing or LogDensityProblems models. +An explicit `SecondOrder` bypasses the gradient slot. Use it for an AD-derived HVP on LogDensityProblems models. See [GPU Execution](15-gpu.md) for backend compatibility. @@ -103,59 +103,21 @@ The prefix-scan kernel runs as pure array broadcasts and is array-type-agnostic. 2. Pass for example, `backend = ADTypes.AutoEnzyme()` to `ParallelMALASampler`. 3. Pass a `CuVector` as `initial_params` to `sample`. -GPU execution has its own caveats (Turing models are CPU-only, `AutoEnzyme()` currently needs the `pmcmc_matmul` / `pmcmc_dot` / `pmcmc_dotsum` wrappers, gradients should avoid scalar indexing). See [GPU Execution](15-gpu.md) for the worked example and the full set of limitations. +GPU execution has its own caveats (`AutoEnzyme()` currently needs the `pmcmc_matmul` / `pmcmc_dot` / `pmcmc_dotsum` wrappers, gradients should avoid scalar indexing). See [GPU Execution](15-gpu.md) for the worked example and the full set of limitations. --- -## Turing.jl integration +## LogDensityProblems models -ParallelMCMC.jl integrates with Turing.jl models through the `DynamicPPL` and `LogDensityProblems` extensions. - -### One-step convenience constructor - -Load `DynamicPPL` (part of Turing.jl) and a single-argument `DensityModel` constructor becomes available: +With `LogDensityProblems` loaded, `DensityModel` accepts any object that implements the [LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface with gradients: ```julia -using Turing, ParallelMCMC, FlexiChains - -@model function normal_model(y) - μ ~ Normal(0.0, 1.0) - y ~ Normal(μ, 0.5) -end - -model = DensityModel(normal_model(1.5)) # param_names=[:μ] extracted automatically +using LogDensityProblems, ParallelMCMC, FlexiChains -chain = sample(model, ParallelMALASampler(0.1; T=64, backend=AutoEnzyme()), 500; - chain_type=VNChain) +model = DensityModel(ld; param_names=[:μ]) # ld is your LogDensityProblems object ``` -Much like Turing's own samplers, the resulting chain will always have parameters in the original (possibly constrained) space, even though the MCMC sampling itself is performed in unconstrained space. -Furthermore, parameter names are automatically extracted from the Turing model (and will always be the same as those when using Turing's own samplers). - -Note that when sampling with a Turing model the returned chain will have `:logjoint`, `:logprior`, and `:loglikelihood` columns, since Turing models provide enough information to separate these contributions to the log-density. -This is in contrast to sampling with a manually constructed `DensityModel`, which returns only a single `:logp` column. - -### Manual `LogDensityProblems` path - -For explicit control over the AD backend, construct the `LogDensityFunction` -directly with DynamicPPL's `adtype` interface: - -```julia -using Turing, LogDensityProblems, ADTypes -using ParallelMCMC, FlexiChains - -ld = DynamicPPL.LogDensityFunction( - normal_model(1.5), - DynamicPPL.getlogjoint_internal, - DynamicPPL.LinkAll(); - adtype=ADTypes.AutoEnzyme(), -) - -model = DensityModel(ld; param_names=[:μ]) -``` - -This also accepts any other `LogDensityProblems`-compatible object. -As above, the returned chain will always contain parameters in the original space: the use of `LinkAll()` above only stipulates that MCMC sampling itself is to be performed in unconstrained space, and has no result on the form of the returned chain. +The object must implement `capabilities` (at least `LogDensityOrder{1}`), `dimension`, `logdensity`, and `logdensity_and_gradient`. `hvp` and the batched slots are forwarded to `DensityModel`; for an AD-derived HVP, pass an explicit `SecondOrder`. --- @@ -207,8 +169,6 @@ model = DensityModel(...; param_names=[:x1, :x2, :x3]) model = DensityModel(...; param_names=[:x, (:y, (1, 2))]) ``` -For Turing.jl models, parameter names are automatically derived from the model and do not need to be specified manually. - --- ## Sequential MALA baselines diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index f3e91fa..baf920d 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -22,13 +22,7 @@ A useful heuristic: if your CPU `grad_logdensity_batch` already saturates BLAS a ## Limitations -### 1. DynamicPPL / Turing models are CPU-only - -The one-argument `DensityModel(::DynamicPPL.Model)` constructor (and the `LogDensityProblems` path that backs it) executes the model with `Vector{Float64}` parameters on CPU. There is no path today that runs a `@model` against `CuArray` parameters. - -To run on GPU, hand-write `logdensity` / `grad_logdensity` (and the batched variants) directly over `CuArray`s and pass them to `DensityModel`. This is shown in the example below. - -### 2. Enzyme on GPU currently needs `pmcmc_matmul` / `pmcmc_dot` / `pmcmc_dotsum` +### 1. Enzyme on GPU currently needs `pmcmc_matmul` / `pmcmc_dot` / `pmcmc_dotsum` `ParallelMCMC` exports three thin wrappers: @@ -58,9 +52,9 @@ When you do **not** need them: These wrappers are a workaround pending upstream Enzyme support for GPU matmul rules. Once that lands, plain `A * B` on `CuArray`s will work and the wrappers will be deprecated. -### 3. Some Enzyme GPU gradients hit a `gc-transition` abort. Staging them as single ops is a workaround +### 2. Some Enzyme GPU gradients hit a `gc-transition` abort. Staging them as single ops is a workaround -Certain `CuArray` gradient expressions abort during Enzyme compilation on GPU with the same `gc-transition` failure as limitation 2. In practice the fused form below aborts while the staged form compiles, so **writing the gradient as single-op stages is a reliable workaround**: +Certain `CuArray` gradient expressions abort during Enzyme compilation on GPU with the same `gc-transition` failure as limitation 1. In practice the fused form below aborts while the staged form compiles, so **writing the gradient as single-op stages is a reliable workaround**: ```julia # ABORTS during Enzyme compile on GPU @@ -253,7 +247,7 @@ sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) **Reactant chooses the execution device independently of the input array.** With a CUDA client, `CuArray` inputs are staged through pinned host memory and results are copied back on the device. With the CPU client, both directions go through the host and the sampler warns when it prepares the model. Select a GPU client with `Reactant.set_default_backend` when available. -- The log-density must be Reactant-traceable. DynamicPPL-built log-densities are not. +- The log-density must be Reactant-traceable. - Across an HVP's two AD passes, use `AutoReactant()` for both or neither. A hand-written gradient may pair with it. `SecondOrder` cannot contain `AutoReactant()`. - Only the default `AutoReactant()` mode is supported. @@ -268,7 +262,7 @@ sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) - **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. +- **You're hitting AD-backend-specific GPU restrictions.** The [Enzyme limitations](#1-enzyme-on-gpu-currently-needs-pmcmc_matmul-pmcmc_dot-pmcmc_dotsum) above (`pmcmc_*` wrappers, staged broadcasts) only matter when the AD backend is invoked. Supplying analytical HVP sidesteps them — your `gradlogp` and `hvp` can use plain `*`, `dot`, `sum`, and the sampler's `backend` can be omitted entirely because no AD is ever invoked. - **You can reuse intermediates between gradient and HVP.** When `hvp` shares $X\beta$, $\sigma(X\beta)$, or similar with the gradient computation, an analytical version can be both faster *and* shorter than what AD produces. ### Same example with analytical HVP diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index ce39e28..9f0f5e7 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -22,12 +22,7 @@ ParallelMCMC.needs_host_staging ## Extension constructors -`DensityModel` also has extension constructors for common probabilistic-programming interfaces: - -- `DensityModel(ld; param_names=nothing)` for `LogDensityProblems` models with gradients -- `DensityModel(turing_model)` for `DynamicPPL` / Turing models when the relevant extension packages are loaded - -See [Getting Started](10-getting-started.md) for end-to-end examples of both. +`DensityModel(ld; param_names=nothing)` wraps any `LogDensityProblems` object with gradients when `LogDensityProblems` is loaded. See [Getting Started](10-getting-started.md) for an example. ## Model diff --git a/docs/src/index.md b/docs/src/index.md index 68aec97..0bd8784 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -95,31 +95,8 @@ chain = sample(model, sampler, 2_000; progress=true) ``` -### Turing.jl integration - -When `DynamicPPL` (part of Turing.jl) is loaded, a one-argument `DensityModel` constructor is available that wraps a `@model` directly. -Parameter names are automatically extracted, and values transformed back to the original model space: - -```julia -using Turing, ParallelMCMC, FlexiChains - -@model function normal_model(y) - μ ~ Normal(0.0, 1.0) - y ~ Normal(μ, 0.5) -end - -model = DensityModel(normal_model(1.5)) -sampler = AdaptiveMALASampler(0.3; n_warmup=500) - -chain = sample(model, sampler, 2_000; - chain_type=VNChain, - discard_warmup=true) -``` - See [Getting Started](10-getting-started.md) for worked examples and guidance on choosing samplers, and [Algorithm Details](20-algorithms.md) for the mathematics behind DEER. -For Turing models the chain type used must be `VNChain` (not `SymChain`), as that is the natural parameter type for Turing models. - ## Contributors ```@raw html diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl deleted file mode 100644 index 53b0599..0000000 --- a/ext/DynamicPPLExt.jl +++ /dev/null @@ -1,167 +0,0 @@ -module DynamicPPLExt - -using ParallelMCMC -using ADTypes: ADTypes -using DynamicPPL: DynamicPPL -using AbstractMCMC: AbstractMCMC -using FlexiChains: FlexiChain, VarName, VNChain, SymChain -using LogDensityProblems: LogDensityProblems - -""" - 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 -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` configures DynamicPPL's `LogDensityFunction`. The remaining -keywords are forwarded to `DensityModel`. HVPs require a callable or explicit -`DifferentiationInterface.SecondOrder`; plain HVP backends are unsupported. -Batched derivatives require a hand-written `logdensity_batch`. - -# Example -```julia -using Turing, ParallelMCMC, FlexiChains, ForwardDiff - -@model function mymodel(y) - μ ~ Normal(0, 1) - y ~ Normal(μ, 0.5) -end - -model = DensityModel(mymodel(1.5); ad_backend=AutoForwardDiff()) -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, - logdensity_batch=nothing, - grad_logdensity_batch=nothing, - hvp_batch=nothing, -) - ld = DynamicPPL.LogDensityFunction( - turing_model, - DynamicPPL.getlogjoint_internal, - DynamicPPL.LinkAll(); - adtype=ad_backend, - ) - return ParallelMCMC.DensityModel( - ld; - hvp=hvp, - logdensity_batch=logdensity_batch, - grad_logdensity_batch=grad_logdensity_batch, - hvp_batch=hvp_batch, - ) -end - -###################### -# Chain construction # -###################### -# In this section, we define overloads for DynamicPPL-based models so that resulting chains -# are converted back into the original parameter space and contain the correct parameter -# names. This is done by converting the raw samples (vectors of parameters) back into -# `DynamicPPL.ParamsWithStats` objects. - -const ParallelMCMCTransitionTypes = Union{ - ParallelMCMC.MALATransition, - ParallelMCMC.AdaptiveMALATransition, - ParallelMCMC.ParallelMALATransition, -} -# Types that represent LogDensityProblems objects that wrap DynamicPPL models. -const LDFPrimal = ParallelMCMC.LogDensityProblemPrimal{<:DynamicPPL.LogDensityFunction} -const LDFGradient = ParallelMCMC.LogDensityProblemGradient{<:DynamicPPL.LogDensityFunction} -const DensityModelLDF = ParallelMCMC.DensityModel{<:LDFPrimal,<:LDFGradient} - -""" - getstats(sample::ParallelMCMCTransitionTypes) - -Get a `NamedTuple` of stats from an MCMC transition. -""" -getstats(sample::ParallelMCMC.MALATransition) = (accepted=sample.accepted,) -function getstats(sample::ParallelMCMC.AdaptiveMALATransition) - return ( - accepted=sample.accepted, step_size=sample.step_size, is_warmup=sample.is_warmup - ) -end -getstats(::ParallelMCMCTransitionTypes) = (;) - -""" - is_warmup(sample::ParallelMCMCTransitionTypes) - -Check if a sample is from the warmup phase of MCMC sampling. -""" -is_warmup(::ParallelMCMCTransitionTypes) = false -is_warmup(sample::ParallelMCMC.AdaptiveMALATransition) = sample.is_warmup - -for (Ttrans, Tspl, Tstate) in ( - (ParallelMCMC.MALATransition, ParallelMCMC.MALASampler, ParallelMCMC.MALAState), - ( - ParallelMCMC.ParallelMALATransition, - ParallelMCMC.ParallelMALASampler, - ParallelMCMC.ParallelMALAState, - ), - ( - ParallelMCMC.AdaptiveMALATransition, - ParallelMCMC.AdaptiveMALASampler, - ParallelMCMC.AdaptiveMALAState, - ), -) - @eval begin - function AbstractMCMC.bundle_samples( - ts::Vector{<:$Ttrans}, - model::DensityModelLDF, - spl::$Tspl, - state::$Tstate, - chain_type::Type{VNChain}; - discard_warmup::Bool=false, - kwargs..., - ) - ts = discard_warmup ? filter(t -> !is_warmup(t), ts) : ts - pwss = map(ts) do t - # Note: This assumes that there is always a field called t.x. This is currently true - # of all samplers in ParallelMCMC - DynamicPPL.ParamsWithStats(t.x, model.logdensity.ld, getstats(t)) - end - return AbstractMCMC.from_samples(VNChain, hcat(pwss)) - end - end - - @eval begin - function AbstractMCMC.bundle_samples( - ts::Vector{<:$Ttrans}, - model::DensityModelLDF, - spl::$Tspl, - state::$Tstate, - chain_type::Type{SymChain}; - kwargs..., - ) - throw( - ArgumentError( - "FlexiChains.SymChain is not supported for DynamicPPL models; please use VNChain instead.", - ), - ) - end - end -end - -function ParallelMCMC._construct_flexichain( - ::Type{VarName}, - vals::AbstractMatrix{<:Real}, - internals::NamedTuple, - ::Any, - model::DensityModelLDF, -) - pwss = map(enumerate(eachrow(vals))) do (i, val) - stats = map(v -> v[i], internals) - return DynamicPPL.ParamsWithStats(val, model.logdensity.ld, stats) - end - return AbstractMCMC.from_samples(VNChain, hcat(pwss)) -end - -end # module diff --git a/ext/LogDensityProblemsExt.jl b/ext/LogDensityProblemsExt.jl index 10f3167..fa771cd 100644 --- a/ext/LogDensityProblemsExt.jl +++ b/ext/LogDensityProblemsExt.jl @@ -14,6 +14,7 @@ Construct a `DensityModel` from any object implementing the - `LogDensityProblems.capabilities(ld)` returning at least `LogDensityProblems.LogDensityOrder{1}` (i.e. gradient available). - `LogDensityProblems.dimension(ld)` -> `Int` +- `LogDensityProblems.logdensity(ld, x)` -> `logp` - `LogDensityProblems.logdensity_and_gradient(ld, x)` -> `(logp, grad)` The optional `param_names` keyword accepts a collection of parameter names that will be used @@ -25,30 +26,25 @@ explicit `DifferentiationInterface.SecondOrder` for HVPs; plain backends and `AutoReactant()` are unsupported. Batched derivatives require a hand-written `logdensity_batch`. -# Turing.jl / DynamicPPL example +# Example ```julia -using Turing, LogDensityProblems, ADTypes, Enzyme, ParallelMCMC, FlexiChains +using LogDensityProblems, ParallelMCMC, FlexiChains -@model function mymodel(y) - μ ~ Normal(0, 1) - y ~ Normal(μ, 0.5) +struct NormalPosterior + y::Float64 +end +LogDensityProblems.capabilities(::Type{NormalPosterior}) = + LogDensityProblems.LogDensityOrder{1}() +LogDensityProblems.dimension(::NormalPosterior) = 1 +LogDensityProblems.logdensity(ld::NormalPosterior, x) = -0.5 * x[1]^2 - 2 * (ld.y - x[1])^2 +function LogDensityProblems.logdensity_and_gradient(ld::NormalPosterior, x) + return LogDensityProblems.logdensity(ld, x), [-x[1] + 4 * (ld.y - x[1])] end -obs = 1.5 -ld = DynamicPPL.LogDensityFunction( - mymodel(obs), - DynamicPPL.getlogjoint_internal, - DynamicPPL.LinkAll(); - adtype=ADTypes.AutoEnzyme(), -) - -model = DensityModel(ld; param_names=[:μ]) +model = DensityModel(NormalPosterior(1.5); param_names=[:μ]) chain = sample(model, AdaptiveMALASampler(0.3; n_warmup=500), 2_000; chain_type=VNChain, discard_warmup=true, progress=true) ``` - -If DynamicPPL is loaded, the simpler one-step constructor `DensityModel(mymodel(obs))` -is also available and extracts parameter names automatically. """ function ParallelMCMC.DensityModel( ld; diff --git a/src/interface.jl b/src/interface.jl index dd68c8e..582f276 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -927,8 +927,7 @@ function _parallel_mala_run!( ) progress = _parallel_mala_progress(progress, progressname) x0 = _parallel_mala_initial_x(rng, model, sampler, initial_params) - # Postprocessing dispatches on the user's model type (e.g. from the - # DynamicPPL extension), so `model` itself stays unprepped. + # Postprocessing receives the user's `model`, not `prepped`. prepped = _prepare_model(model, x0, sampler.T, sampler.backend) ws = nothing nsteps = 0 diff --git a/test/Project.toml b/test/Project.toml index a2370e2..183622d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,8 +2,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" diff --git a/test/test-DEER-Turing-Logistic.jl b/test/test-DEER-Logistic.jl similarity index 70% rename from test/test-DEER-Turing-Logistic.jl rename to test/test-DEER-Logistic.jl index bc5b20f..939294b 100644 --- a/test/test-DEER-Turing-Logistic.jl +++ b/test/test-DEER-Logistic.jl @@ -6,12 +6,8 @@ using FlexiChains using ParallelMCMC -using DynamicPPL -using LogDensityProblems using ADTypes using Enzyme -using ForwardDiff -using Distributions: MvNormal, Bernoulli #= Bayesian logistic regression: @@ -35,8 +31,8 @@ end const _LR_X, _LR_y = _lr_data(MersenneTwister(1234), _LR_N, _LR_D, _LR_β_true) #= -Array-type-agnostic logp and gradlogp — identical math as the Turing model, -but expressed as plain broadcasts so they run on CPU or GPU without change. +Array-type-agnostic logp and gradlogp, written as plain broadcasts so they run +on CPU or GPU without change. =# function _logp_lr(β, X, y) logits = X * β @@ -57,94 +53,6 @@ function _hvp_lr(β, v, X, y) return -(X' * (w .* (X * v))) .- v end -@model function _deer_logistic_regression(X, y) - D = size(X, 2) - β ~ MvNormal(zeros(D), I) - for i in eachindex(y) - p_i = 1 / (1 + exp(-dot(view(X, i, :), β))) - y[i] ~ Bernoulli(p_i) - end -end - -function _deer_logistic_turing_density_model() - return DensityModel( - _deer_logistic_regression(_LR_X, _LR_y); - ad_backend=ADTypes.AutoForwardDiff(), - hvp=(β, v) -> _hvp_lr(β, v, _LR_X, _LR_y), - ) -end - -# Turing integration tests (CPU) - -@testset "ParallelMALASampler Turing logistic: param names extracted correctly" begin - model = _deer_logistic_turing_density_model() - - @test model.dim == _LR_D - @test isfinite(model.logdensity(zeros(_LR_D))) - @test all(isfinite, model.grad_logdensity(zeros(_LR_D))) -end - -@testset "ParallelMALASampler Turing logistic: chains output well-formed" begin - model = _deer_logistic_turing_density_model() - sampler = ParallelMALASampler( - 0.05; - T=16, - maxiter=80, - tol_abs=1e-4, - tol_rel=1e-3, - jacobian=:diag, - damping=0.5, - backend=ADTypes.AutoEnzyme(), - ) - - chain = sample( - MersenneTwister(42), - model, - sampler, - 400; - initial_params=zeros(_LR_D), - chain_type=VNChain, - progress=false, - ) - - @test chain isa VNChain - @test FlexiChains.niters(chain) == 400 - @test @varname(β) in FlexiChains.parameters(chain) - @test all(isfinite, chain[@varname(β), stack = true]) -end - -@testset "ParallelMALASampler Turing logistic: posterior sign correct" begin - model = _deer_logistic_turing_density_model() - sampler = ParallelMALASampler( - 0.05; - T=16, - maxiter=80, - tol_abs=1e-4, - tol_rel=1e-3, - jacobian=:diag, - damping=0.5, - backend=ADTypes.AutoEnzyme(), - ) - - chain = sample( - MersenneTwister(99), - model, - sampler, - 800; - initial_params=zeros(_LR_D), - chain_type=VNChain, - progress=false, - ) - - post = chain[@varname(β), stack = true][201:end, :, :] - β_mean = vec(mean(post; dims=1)) - - @test sign(β_mean[1]) == sign(_LR_β_true[1]) - @test sign(β_mean[2]) == sign(_LR_β_true[2]) - @test abs(β_mean[1] - _LR_β_true[1]) < 0.6 - @test abs(β_mean[2] - _LR_β_true[2]) < 0.6 -end - # Statistical correctness: DEER posterior mean matches AdaptiveMALA reference. @testset "ParallelMALASampler logistic: posterior mean matches AdaptiveMALA reference" begin diff --git a/test/test-GPU-AD-HVP.jl b/test/test-GPU-AD-HVP.jl index 0d0a8c4..e841ec2 100644 --- a/test/test-GPU-AD-HVP.jl +++ b/test/test-GPU-AD-HVP.jl @@ -11,7 +11,7 @@ using Zygote using CUDA: CUDA #= -Coverage for the GPU AD-HVP path: when the user supplies a non-Turing model +Coverage for the GPU AD-HVP path: when the user supplies a model with `gradlogp` / `grad_logdensity_batch` but no analytical `hvp`/`hvp_batch`, the sampler should fall back to AD on the user's gradient and run on GPU without crashing. diff --git a/test/test-LogDensityProblems.jl b/test/test-LogDensityProblems.jl new file mode 100644 index 0000000..bea1d12 --- /dev/null +++ b/test/test-LogDensityProblems.jl @@ -0,0 +1,81 @@ +using Test +using Random +using Statistics +using FlexiChains + +using ParallelMCMC + +using LogDensityProblems: LogDensityProblems +using ADTypes +using ForwardDiff: ForwardDiff + +#= +μ ~ N(0, 1), y | μ ~ N(μ, 0.5), observed y = 1.5. +Posterior: σ_post² = 1 / (1 + 1/0.25) = 0.2, μ_post = 0.2 * 1.5 / 0.25 = 1.2. +The Hessian is constant: -1 - 1/0.25 = -5. +=# +const LDP_OBS = 1.5 +const LDP_MU_POST = 1.2 +const LDP_VAR_POST = 0.2 + +struct _NormalLD end +LogDensityProblems.capabilities(::Type{_NormalLD}) = LogDensityProblems.LogDensityOrder{1}() +LogDensityProblems.dimension(::_NormalLD) = 1 +function LogDensityProblems.logdensity(::_NormalLD, x) + μ = x[1] + return -0.5 * μ^2 - 0.5 * ((LDP_OBS - μ) / 0.5)^2 +end +function LogDensityProblems.logdensity_and_gradient(ld::_NormalLD, x) + μ = x[1] + return LogDensityProblems.logdensity(ld, x), [-μ + (LDP_OBS - μ) / 0.25] +end + +struct _NoGradLD end +LogDensityProblems.capabilities(::Type{_NoGradLD}) = LogDensityProblems.LogDensityOrder{0}() +LogDensityProblems.dimension(::_NoGradLD) = 1 + +@testset "constructor" begin + model = DensityModel(_NormalLD()) + @test model.dim == 1 + @test model.logdensity([0.0]) ≈ -4.5 + @test model.grad_logdensity([0.0]) ≈ [6.0] + @test model.param_names === nothing + + @test_throws "must support gradients" DensityModel(_NoGradLD()) +end + +@testset "AdaptiveMALA recovers the posterior" begin + model = DensityModel(_NormalLD(); param_names=[:μ]) + chain = sample( + MersenneTwister(1), + model, + AdaptiveMALASampler(0.3; n_warmup=500), + 4_000; + chain_type=VNChain, + discard_warmup=true, + progress=false, + ) + @test only(FlexiChains.parameters(chain)) == @varname(μ) + draws = vec(chain[@varname(μ)]) + @test abs(mean(draws) - LDP_MU_POST) < 0.1 + @test abs(var(draws) - LDP_VAR_POST) < 0.1 +end + +@testset "SecondOrder hvp drives ParallelMALA" begin + so = ParallelMCMC.DI.SecondOrder(ADTypes.AutoForwardDiff(), ADTypes.AutoForwardDiff()) + model = DensityModel(_NormalLD(); param_names=[:μ], hvp=so) + @test model.hvp === so + + prepped = ParallelMCMC._prepare_model(model, [0.0], 8, nothing) + @test prepped.hvp([0.0], [1.0]) ≈ [-5.0] + + chain = sample( + MersenneTwister(12), + model, + ParallelMALASampler(0.02; T=8), + 64; + chain_type=SymChain, + progress=false, + ) + @test all(isfinite, vec(chain[:μ])) +end diff --git a/test/test-Turing-Integration.jl b/test/test-Turing-Integration.jl deleted file mode 100644 index 501b5ec..0000000 --- a/test/test-Turing-Integration.jl +++ /dev/null @@ -1,384 +0,0 @@ -using Test -using Random -using LinearAlgebra -using Statistics -using FlexiChains - -using ParallelMCMC - -using DynamicPPL -using LogDensityProblems -using ADTypes -using Enzyme -using ForwardDiff -using Distributions: Beta, Dirichlet, Normal, MvNormal, product_distribution - -#= -A simple 1-D normal likelihood: μ ~ N(0,1), y | μ ~ N(μ, 0.5) -Posterior: μ | y=1.5 is N(μ_post, σ_post²) -σ_post² = 1 / (1/1² + 1/0.5²) = 1 / (1 + 4) = 0.2 -μ_post = σ_post² * (y / 0.5²) = 0.2 * (1.5 / 0.25) = 0.2 * 6 = 1.2 -=# -const TRUE_OBS = 1.5 -const TRUE_MU_POST = 1.2 -const TRUE_VAR_POST = 0.2 - -@model function normal_model(y) - μ ~ Normal(0.0, 1.0) - y ~ Normal(μ, 0.5) -end - -@model function mv_model(y) - c ~ Dirichlet(ones(3)) # to test constraints - μ ~ product_distribution((a=Normal(), b=Normal())) - y ~ MvNormal([μ.a, μ.b], 0.5 * I) -end - -@model function beta_model() - x ~ Beta(2, 2) -end - -@model function mvnormal_2d_model() - x ~ MvNormal(zeros(2), I) -end - -@model function dirichlet_3_model() - x ~ Dirichlet(ones(3)) -end - -@testset "directly passing LogDensityFunction" begin - ld = DynamicPPL.LogDensityFunction( - normal_model(TRUE_OBS), - DynamicPPL.getlogjoint_internal, - DynamicPPL.LinkAll(); - adtype=ADTypes.AutoEnzyme(), - ) - - model = DensityModel(ld) - - @test model.dim == 1 - - chain = sample( - MersenneTwister(1), - model, - AdaptiveMALASampler(0.3; n_warmup=200), - 600; - chain_type=VNChain, - progress=false, - ) - @test only(FlexiChains.parameters(chain)) == @varname(μ) -end - -@testset "DynamicPPLExt: convenience constructor" begin - model = DensityModel(normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff()) - - @test model.dim == 1 - @test isfinite(model.logdensity([0.0])) - @test isfinite(model.grad_logdensity([0.0])[1]) - - chain = sample( - MersenneTwister(1), - model, - AdaptiveMALASampler(0.3; n_warmup=200), - 600; - chain_type=VNChain, - progress=false, - ) - @test only(FlexiChains.parameters(chain)) == @varname(μ) - - @test_throws "SymChain is not supported" sample( - MersenneTwister(1), - model, - AdaptiveMALASampler(0.3; n_warmup=200), - 600; - chain_type=SymChain, - progress=false, - ) -end - -@testset "DynamicPPLExt: convenience constructor uses linked space for constrained models" begin - model = DensityModel(beta_model(); ad_backend=ADTypes.AutoForwardDiff()) - - @test model.dim == 1 - @test isfinite(model.logdensity([-0.4])) - @test isfinite(model.grad_logdensity([-0.4])[1]) -end - -@testset "DynamicPPLExt: generic Turing model works with ParallelMALA and default Enzyme HVP" begin - model = DensityModel(normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff()) - - @test model.hvp === nothing - - for jacobian in (:diag, :stoch_diag) - sampler = ParallelMALASampler( - 0.03; - T=4, - maxiter=80, - tol_abs=1e-5, - tol_rel=1e-4, - jacobian=jacobian, - damping=0.5, - backend=ADTypes.AutoEnzyme(), - ) - - trans, state = ParallelMCMC.AbstractMCMC.step( - MersenneTwister(11), model, sampler; initial_params=[0.0] - ) - - @test trans isa ParallelMALATransition - @test state isa ParallelMALAState - @test isfinite(trans.logp) - @test all(isfinite, state.trajectory) - end -end - -@testset "DynamicPPLExt: SecondOrder hvp on a Turing model" begin - # For this model, H = -1 - 1/0.5^2 = -5. - so = ParallelMCMC.DI.SecondOrder(ADTypes.AutoForwardDiff(), ADTypes.AutoForwardDiff()) - model = DensityModel( - normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff(), hvp=so - ) - @test model.hvp === so - - prepped = ParallelMCMC._prepare_model(model, [0.0], 8, nothing) - @test prepped.hvp([0.0], [1.0]) ≈ [-5.0] - chain = sample( - MersenneTwister(12), - model, - ParallelMALASampler(0.02; T=8), - 64; - chain_type=VNChain, - progress=false, - ) - @test all(isfinite, vec(chain[@varname(μ)])) - - model_plain = DensityModel( - normal_model(TRUE_OBS); - ad_backend=ADTypes.AutoForwardDiff(), - hvp=ADTypes.AutoForwardDiff(), - ) - prepped_plain = ParallelMCMC._prepare_model(model_plain, [0.0], 8, nothing) - @test_throws Exception prepped_plain.hvp([0.0], [1.0]) -end - -@testset "DynamicPPLExt: batched slots reach the batched DEER path" begin - σ = 0.5 - logp_b(X) = - vec(-0.5 .* X .^ 2 .- 0.5 .* ((TRUE_OBS .- X) ./ σ) .^ 2 .- log(2π) .- log(σ)) - grad_b(X) = -X .+ (TRUE_OBS .- X) ./ σ^2 - hvp_b(X, V) = (-1 - 1 / σ^2) .* V - - model = DensityModel( - normal_model(TRUE_OBS); - ad_backend=ADTypes.AutoForwardDiff(), - hvp=(x, v) -> (-1 - 1 / σ^2) .* v, - logdensity_batch=logp_b, - grad_logdensity_batch=grad_b, - hvp_batch=hvp_b, - ) - - @test model.logdensity_batch === logp_b - @test model.grad_logdensity_batch === grad_b - @test model.hvp_batch === hvp_b - - X = reshape([-0.5, 0.0, 0.7, 1.4], 1, 4) - @test logp_b(X) ≈ [model.logdensity(X[:, t]) for t in 1:size(X, 2)] - - prepped = ParallelMCMC._prepare_model(model, [0.0], 4, nothing) - @test prepped.grad_logdensity_batch === grad_b - @test prepped.hvp_batch === hvp_b - - chain = sample( - MersenneTwister(13), - model, - ParallelMALASampler(0.02; T=8), - 64; - chain_type=VNChain, - progress=false, - ) - @test all(isfinite, vec(chain[@varname(μ)])) -end - -@testset "DynamicPPLExt: MvNormal(zeros(2), I) runs with ParallelMALA" begin - model = DensityModel(mvnormal_2d_model(); ad_backend=ADTypes.AutoForwardDiff()) - - @test model.dim == 2 - @test isfinite(model.logdensity(zeros(2))) - @test all(isfinite, model.grad_logdensity(zeros(2))) - - sampler = ParallelMALASampler( - 0.2; T=8, maxiter=80, tol_abs=1e-4, tol_rel=1e-3, backend=ADTypes.AutoEnzyme() - ) - chain = sample( - MersenneTwister(3), - model, - sampler, - 800; - initial_params=zeros(2), - chain_type=VNChain, - progress=false, - ) - samples = chain[@varname(x), stack = true] - @test all(isfinite, samples) - # Standard normal in 2-D: posterior mean should be near zero. - posterior_means = mean(samples; dims=1) - @test maximum(abs, posterior_means) < 0.25 -end - -@testset "DynamicPPLExt: Dirichlet(ones(3)) runs with ParallelMALA (linked space)" begin - #= - Dirichlet(ones(3)) lives on a 2-simplex, so its unconstrained - representation has dim 2. Bijectors handles the link/unlink. - =# - model = DensityModel(dirichlet_3_model(); ad_backend=ADTypes.AutoForwardDiff()) - - @test model.dim == 2 - @test isfinite(model.logdensity(zeros(2))) - @test all(isfinite, model.grad_logdensity(zeros(2))) - - sampler = ParallelMALASampler( - 0.2; T=8, maxiter=80, tol_abs=1e-4, tol_rel=1e-3, backend=ADTypes.AutoEnzyme() - ) - chain = sample( - MersenneTwister(4), - model, - sampler, - 800; - initial_params=zeros(2), - chain_type=VNChain, - progress=false, - ) - @test chain isa VNChain - @test all(isfinite, Array(chain)) -end - -@testset "DynamicPPLExt: ParallelMALA bundle_samples fallback path (callback)" begin - #= A per-step `callback` forces ParallelMALA's `mcmcsample` override down the - generic step/bundle_samples path =# - model = DensityModel(mvnormal_2d_model(); ad_backend=ADTypes.AutoForwardDiff()) - sampler = ParallelMALASampler( - 0.2; T=8, maxiter=80, tol_abs=1e-4, tol_rel=1e-3, backend=ADTypes.AutoEnzyme() - ) - chain = sample( - MersenneTwister(3), - model, - sampler, - 800; - initial_params=zeros(2), - chain_type=VNChain, - callback=(args...; kwargs...) -> nothing, - progress=false, - ) - @test chain isa VNChain - @test only(FlexiChains.parameters(chain)) == @varname(x) - @test all(isfinite, Array(chain)) -end - -@testset "DynamicPPLExt: ParallelMALA fast path with thinning matches the step-wise loop" begin - model = DensityModel(mvnormal_2d_model(); ad_backend=ADTypes.AutoForwardDiff()) - sampler = ParallelMALASampler( - 0.2; T=8, maxiter=80, tol_abs=1e-4, tol_rel=1e-3, backend=ADTypes.AutoEnzyme() - ) - kw = ( - initial_params=zeros(2), - chain_type=VNChain, - thinning=2, - discard_initial=10, - progress=false, - ) - chain = sample(MersenneTwister(3), model, sampler, 200; kw...) - slow = ParallelMCMC._default_parallel_mala_mcmcsample( - MersenneTwister(3), model, sampler, 200; kw... - ) - @test chain isa VNChain - @test only(FlexiChains.parameters(chain)) == @varname(x) - @test FlexiChains.niters(chain) == 200 - @test all(isfinite, Array(chain)) - @test chain[@varname(x), stack = true] == slow[@varname(x), stack = true] -end - -@testset "DynamicPPLExt: named columns in Chains output" begin - model = DensityModel(normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff()) - - chain = sample( - MersenneTwister(2), - model, - AdaptiveMALASampler(0.3; n_warmup=200), - 600; - chain_type=VNChain, - progress=false, - ) - - @test chain isa VNChain - params = FlexiChains.parameters(chain) - @test @varname(μ) in params - @test !(@varname(x) in params) -end - -@testset "discard_warmup=true removes warmup samples" begin - model = DensityModel(normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff()) - n_warmup = 200 - n_total = 800 - sampler = AdaptiveMALASampler(0.3; n_warmup=n_warmup) - - chain_full = sample( - MersenneTwister(3), model, sampler, n_total; chain_type=VNChain, progress=false - ) - chain_trimmed = sample( - MersenneTwister(3), - model, - sampler, - n_total; - chain_type=VNChain, - progress=false, - discard_warmup=true, - ) - - @test FlexiChains.niters(chain_full) == n_total - @test FlexiChains.niters(chain_trimmed) == n_total - n_warmup - 1 - @test all(==(false), chain_trimmed[:is_warmup]) -end - -@testset "posterior mean and variance match analytic solution" begin - model = DensityModel(normal_model(TRUE_OBS); ad_backend=ADTypes.AutoForwardDiff()) - n_warmup = 2_000 - n_draw = 10_000 - sampler = AdaptiveMALASampler(0.3; n_warmup=n_warmup) - - chain = sample( - MersenneTwister(2025), - model, - sampler, - n_warmup + n_draw; - chain_type=VNChain, - progress=false, - discard_warmup=true, - ) - - mu_samples = vec(chain[:μ]) - - @test abs(mean(mu_samples) - TRUE_MU_POST) < 0.05 - @test abs(var(mu_samples) - TRUE_VAR_POST) < 0.05 -end - -@testset "multivariate model: named columns for each dimension" begin - obs = [1.0, -1.0] - model = DensityModel(mv_model(obs); ad_backend=ADTypes.AutoForwardDiff()) - - # 2 from linked Dirichlet + 2 from product_distribution - @test model.dim == 4 - - chain = sample( - MersenneTwister(7), - model, - AdaptiveMALASampler(0.2; n_warmup=100), - 300; - chain_type=VNChain, - progress=false, - ) - - # Check that the chain contains parameters in original space. - # The Dirichlet parameter should have length 3. - @test Set(FlexiChains.parameters(chain)) == Set([@varname(c), @varname(μ)]) - @test all(chain[@varname(c), stack = true] .>= 0.0) # Dirichlet samples should be non-negative -end From 594d8987bbbd704f6bc2ae0cdcc45085dc2b3199 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:50:53 -0400 Subject: [PATCH 2/2] Address review on LogDensityProblems docs and tests - Reword the comment on the LogDensityProblems wrapper structs; their remaining dispatch use is the Reactant HVP check. Drop a comment that only restated code. - A plain backend in `hvp` works when `logdensity_and_gradient` is differentiable by it. Say that instead of calling plain backends unsupported, and keep `SecondOrder` as the fallback. - Put a runnable LogDensityProblems example in Getting Started and link it from the reference page. - Test that the batched slots pass through the LogDensityProblems constructor, that the default parameter name is `x`, and tighten the variance tolerance. --- docs/src/10-getting-started.md | 22 ++++++++++++++--- docs/src/95-reference.md | 4 +-- ext/LogDensityProblemsExt.jl | 8 +++--- src/interface.jl | 8 +++--- test/test-LogDensityProblems.jl | 43 ++++++++++++++++++++++++++++++++- 5 files changed, 71 insertions(+), 14 deletions(-) diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index 56269d2..7475b75 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -47,7 +47,7 @@ Hand-written and AD-derived slots can be mixed. | backend | backend | `SecondOrder(hvp, grad_logdensity)` on `logdensity` | | either | `SecondOrder(...)` | that pair on `logdensity`, gradient slot unused | -An explicit `SecondOrder` bypasses the gradient slot. Use it for an AD-derived HVP on LogDensityProblems models. +An explicit `SecondOrder` bypasses the gradient slot. Use it when the gradient can't be differentiated. See [GPU Execution](15-gpu.md) for backend compatibility. @@ -114,10 +114,26 @@ With `LogDensityProblems` loaded, `DensityModel` accepts any object that impleme ```julia using LogDensityProblems, ParallelMCMC, FlexiChains -model = DensityModel(ld; param_names=[:μ]) # ld is your LogDensityProblems object +# μ ~ Normal(0, 1), y ~ Normal(μ, 0.5) +struct NormalPosterior + y::Float64 +end +LogDensityProblems.capabilities(::Type{NormalPosterior}) = + LogDensityProblems.LogDensityOrder{1}() +LogDensityProblems.dimension(::NormalPosterior) = 1 +LogDensityProblems.logdensity(ld::NormalPosterior, x) = -0.5 * x[1]^2 - 2 * (ld.y - x[1])^2 +function LogDensityProblems.logdensity_and_gradient(ld::NormalPosterior, x) + return LogDensityProblems.logdensity(ld, x), [-x[1] + 4 * (ld.y - x[1])] +end + +model = DensityModel(NormalPosterior(1.5); param_names=[:μ]) +chain = sample(model, AdaptiveMALASampler(0.3; n_warmup=500), 2_000; + chain_type=VNChain, discard_warmup=true) ``` -The object must implement `capabilities` (at least `LogDensityOrder{1}`), `dimension`, `logdensity`, and `logdensity_and_gradient`. `hvp` and the batched slots are forwarded to `DensityModel`; for an AD-derived HVP, pass an explicit `SecondOrder`. +The object must implement `capabilities` (at least `LogDensityOrder{1}`), `dimension`, `logdensity`, and `logdensity_and_gradient`. `hvp` and the batched slots are forwarded to `DensityModel`. + +A backend in `hvp` differentiates `logdensity_and_gradient`, so that function must be differentiable by the backend. If it is not (for example, it calls AD internally), pass an explicit `SecondOrder`, which differentiates `logdensity` twice instead. `AutoReactant()` is not supported for these models. --- diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index 9f0f5e7..167a0fc 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -20,9 +20,9 @@ derivative-slot backend on `DensityModel` and as `ParallelMALASampler`'s ParallelMCMC.needs_host_staging ``` -## Extension constructors +## Extension constructor -`DensityModel(ld; param_names=nothing)` wraps any `LogDensityProblems` object with gradients when `LogDensityProblems` is loaded. See [Getting Started](10-getting-started.md) for an example. +`DensityModel(ld; param_names=nothing)` wraps any `LogDensityProblems` object with gradients when `LogDensityProblems` is loaded. See [LogDensityProblems models](@ref) for an example. ## Model diff --git a/ext/LogDensityProblemsExt.jl b/ext/LogDensityProblemsExt.jl index fa771cd..9fc94cc 100644 --- a/ext/LogDensityProblemsExt.jl +++ b/ext/LogDensityProblemsExt.jl @@ -21,9 +21,11 @@ The optional `param_names` keyword accepts a collection of parameter names that for the columns of the returned `FlexiChain` object. If omitted, a single vector-valued parameter named `:x` will be chosen, unless you also pass `param_names` to `sample(...)`. -`hvp` and the batched slots are forwarded to `DensityModel`. Use a callable or -explicit `DifferentiationInterface.SecondOrder` for HVPs; plain backends and -`AutoReactant()` are unsupported. Batched derivatives require a hand-written +`hvp` and the batched slots are forwarded to `DensityModel`. A backend in `hvp` +differentiates `logdensity_and_gradient`, so that function must be differentiable by the +backend. If it is not (for example, it calls AD internally), pass an explicit +`DifferentiationInterface.SecondOrder`, which differentiates `logdensity` twice instead. +`AutoReactant()` is unsupported. Batched derivatives require a hand-written `logdensity_batch`. # Example diff --git a/src/interface.jl b/src/interface.jl index 582f276..8935f84 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -392,10 +392,9 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, ) end -#= Callable structs that allow us to dispatch on the type of the LogDensityProblems object in -the postprocessing stage. Ideally these would be defined in the LogDensityProblemsExt. -However, structs defined in extensions are hard to get hold of so we define them here. -The callable behaviour itself is implemented in LogDensityProblemsExt =# +#= Wrappers for LogDensityProblems objects. They are defined here rather than in +LogDensityProblemsExt so `_check_reactant_hvp_source` can dispatch on them. The call +methods live in the extension. =# struct LogDensityProblemPrimal{L} ld::L end @@ -927,7 +926,6 @@ function _parallel_mala_run!( ) progress = _parallel_mala_progress(progress, progressname) x0 = _parallel_mala_initial_x(rng, model, sampler, initial_params) - # Postprocessing receives the user's `model`, not `prepped`. prepped = _prepare_model(model, x0, sampler.T, sampler.backend) ws = nothing nsteps = 0 diff --git a/test/test-LogDensityProblems.jl b/test/test-LogDensityProblems.jl index bea1d12..89d9542 100644 --- a/test/test-LogDensityProblems.jl +++ b/test/test-LogDensityProblems.jl @@ -58,7 +58,19 @@ end @test only(FlexiChains.parameters(chain)) == @varname(μ) draws = vec(chain[@varname(μ)]) @test abs(mean(draws) - LDP_MU_POST) < 0.1 - @test abs(var(draws) - LDP_VAR_POST) < 0.1 + @test abs(var(draws) - LDP_VAR_POST) < 0.05 +end + +@testset "parameter is named x by default" begin + chain = sample( + MersenneTwister(2), + DensityModel(_NormalLD()), + MALASampler(0.3), + 20; + chain_type=SymChain, + progress=false, + ) + @test only(FlexiChains.parameters(chain)) == :x end @testset "SecondOrder hvp drives ParallelMALA" begin @@ -79,3 +91,32 @@ end ) @test all(isfinite, vec(chain[:μ])) end + +@testset "batched slots reach the prepared model" begin + logp_b(X) = vec(-0.5 .* X .^ 2 .- 2 .* (LDP_OBS .- X) .^ 2) + grad_b(X) = -X .+ 4 .* (LDP_OBS .- X) + hvp_b(X, V) = -5 .* V + model = DensityModel( + _NormalLD(); + param_names=[:μ], + hvp=(x, v) -> -5 .* v, + logdensity_batch=logp_b, + grad_logdensity_batch=grad_b, + hvp_batch=hvp_b, + ) + + prepped = ParallelMCMC._prepare_model(model, [0.0], 8, nothing) + @test prepped.logdensity_batch === logp_b + @test prepped.grad_logdensity_batch === grad_b + @test prepped.hvp_batch === hvp_b + + chain = sample( + MersenneTwister(13), + model, + ParallelMALASampler(0.02; T=8), + 64; + chain_type=SymChain, + progress=false, + ) + @test all(isfinite, vec(chain[:μ])) +end