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 (28.68%) 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 @@
## PR-Fixes #66 +/- ##
============================================
- Coverage 92.05% 86.56% -5.49%
============================================
Files 7 9 +2
Lines 1170 1280 +110
============================================
+ Hits 1077 1108 +31
- Misses 93 172 +79 ☔ 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 |
CUDA moves from `[deps]` to `[weakdeps]` behind a new `CUDAExt`, so installing and loading ParallelMCMC no longer drags the CUDA stack onto machines that cannot use it. `using ParallelMCMC` alone now leaves CUDA unloaded. The samplers were already array-type-agnostic everywhere but the random fills -- MALA's normal noise and DEER's Rademacher probes -- which cannot be written into device memory one element at a time. Those four `isa CuArray` checks all asked the same question, so they collapse into one trait, `needs_host_staging(x)`, defaulting to `false`. `_rademacher!` loses its two `CuArray` dispatches for a trait branch in the generic method. The extension is a single method setting that trait for `CuArray`, which also drops the NVIDIA-only assumption: an AMDGPU or Metal user adds one method rather than waiting on another extension. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Resolves #37 for real, though its limited to Reactant for GPU (as of now)