Reactant integration - #66
Conversation
3bea841 to
3e4447b
Compare
There was a problem hiding this comment.
Pull request overview
This PR adds a Reactant-based derivative/HVP path (ADTypes.AutoReactant()) to enable true second-order HVPs via Reactant/XLA (bypassing DifferentiationInterface) and updates docs/tests to cover pairing rules, limitations, and opt-in test execution. It primarily targets unblocking GPU second-order behavior discussed in #37 by providing a non-DI route.
Changes:
- Adds
ReactantExtto compile gradients/HVPs with Reactant (@compile) and introduces load-hint fallbacks when Reactant isn’t loaded. - Adds extensive tests for Reactant pairing rules, error surfacing (pre-compile), analytic correctness, and end-to-end sampling (opt-in).
- Updates docs and changelog to document Reactant behavior, caveats, and GPU guidance.
Reviewed changes
Copilot reviewed 15 out of 16 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| test/test-Reactant-HVP.jl | New opt-in test suite covering Reactant pairing/error behavior and correctness. |
| test/test-HVP-Strategy.jl | Adds strategy-selection test ensuring AutoReactant bypasses DI.hvp_mode. |
| test/runtests.jl | Gates Reactant tests behind PARALLELMCMC_TEST_REACTANT env var. |
| test/Project.toml | Adds Reactant as a test extra (not default-installed). |
| src/ParallelMCMC.jl | Adds shared _REACTANT_LOAD_HINT constant for consistent error messaging. |
| src/interface.jl | Adds Reactant gradient hooks, pairing checks, and second-order backend composition for Reactant. |
| src/DEER/DEER.jl | Adds ReactantHVP strategy + load-hint fallbacks for missing Reactant extension. |
| Project.toml | Registers ReactantExt extension and adds Reactant as a weak dependency with compat. |
| ext/ReactantExt.jl | Implements compiled gradient/HVP factories using Reactant + Enzyme-MLIR tracing. |
| ext/LogDensityProblemsExt.jl | Documents that AutoReactant cannot trace LDP/DynamicPPL machinery. |
| ext/EnzymeExt.jl | Documentation edits clarifying normalization behavior (no mode forcing). |
| ext/DynamicPPLExt.jl | Documentation edits clarifying HVP expectations and Reactant incompatibility. |
| docs/src/95-reference.md | Adds Reactant section to the API reference pointing to GPU docs/caveats. |
| docs/src/15-gpu.md | Documents Reactant path, caveats, pairing rules, and relation to #37. |
| docs/src/10-getting-started.md | Updates backend pairing guidance to include AutoReactant exception on GPU. |
| CHANGELOG.md | Records Reactant integration and related behavior changes/fixes. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report❌ Patch coverage is
❌ Your patch check has failed because the patch coverage (26.82%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #66 +/- ##
==========================================
- Coverage 92.05% 85.69% -6.36%
==========================================
Files 7 8 +1
Lines 1170 1279 +109
==========================================
+ Hits 1077 1096 +19
- Misses 93 183 +90 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…llelMCMC.jl into reactant-integration
|
How's that working out? |
|
So far so good? I've tried it on some mundane models (e.g., MvNormal, LinReg) and am now trying to get it to work for more complicated models (LogReg, Banana shaped Posterior) and currently performance is pretty good so generally very encouraging |
54ec42b to
4dc9a62
Compare
PR #65 landed the #62 review work on main; this branch still carried an earlier draft of the same edits. Conflicts in CHANGELOG.md, ext/EnzymeExt.jl, src/DEER/DEER.jl and test/test-HVP-Strategy.jl resolve to main's text, with the Reactant-only additions re-applied on top: the ReactantHVP strategy and its load-hint fallbacks, the AutoReactant testset, and the Reactant changelog entries.
The merge resolved the shared #62 comment blocks to main's fuller wording; these are the trimmed versions from `cut down chaff`. Comments only — the code is still identical to main apart from the Reactant additions.
|
@wsmoses if you had a moment, would you mind reviewing this? |
| 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))) |
There was a problem hiding this comment.
this is fine to start with but longer term for perf we should avoid having to constantly copy data back/forth
Resolves #37 for real, though its limited to Reactant for GPU (as of now)