Reactant integration

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
 

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"

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