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
Additional details and impacted files@@ Coverage Diff @@
## PR-Fixes #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 |
Resolves #37 for real, though its limited to Reactant for GPU (as of now)