Integrate DifferentiationInterface.jl for gradient computation

[shravanngoswamii]

Replace LogDensityProblemsAD with DifferentiationInterface.jl for automatic differentiation.

Changes:

• Add adtype parameter to compile() for specifying AD backends
• Support symbol shortcuts: :ReverseDiff, :ForwardDiff, :Zygote, :Enzyme
• Implement BUGSModelWithGradient wrapper with logdensity_and_gradient method
• Backward compatible: existing code without adtype continues to work

Usage:

model = compile(model_def, data; adtype=:ReverseDiff)

model = compile(model_def, data; adtype=AutoReverseDiff(compile=true)) # Same as above

Closes #380

[github-actions]

JuliaBUGS.jl documentation for PR #397 is available at:
https://TuringLang.github.io/JuliaBUGS.jl/previews/PR397/

[shravanngoswamii]

Updated Docs: https://turinglang.org/JuliaBUGS.jl/previews/PR397/example/#Automatic-Differentiation

[coveralls]

Pull Request Test Coverage Report for Build 18304377034

Details

• 25 of 45 (55.56%) changed or added relevant lines in 5 files are covered.
• 4 unchanged lines in 2 files lost coverage.
• Overall coverage decreased (-0.4%) to 82.873%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
JuliaBUGS/ext/JuliaBUGSAdvancedHMCExt.jl	4	9	44.44%
JuliaBUGS/src/model/logdensityproblems.jl	10	16	62.5%
JuliaBUGS/ext/JuliaBUGSAdvancedMHExt.jl	0	9	0.0%

Files with Coverage Reduction	New Missed Lines	%
JuliaBUGS/ext/JuliaBUGSAdvancedMHExt.jl	1	35.14%
JuliaBUGS/ext/JuliaBUGSMCMCChainsExt.jl	3	72.0%

Totals
Change from base Build 17946015706:	-0.4%
Covered Lines:	3150
Relevant Lines:	3801

---

💛 - Coveralls

[codecov]

Codecov Report

❌ Patch coverage is 55.55556% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.87%. Comparing base (53fd7bb) to head (ef36d7d).

Files with missing lines	Patch %	Lines
JuliaBUGS/ext/JuliaBUGSAdvancedMHExt.jl	0.00%	9 Missing ⚠️
JuliaBUGS/src/model/logdensityproblems.jl	62.50%	6 Missing ⚠️
JuliaBUGS/ext/JuliaBUGSAdvancedHMCExt.jl	44.44%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #397      +/-   ##
==========================================
- Coverage   83.26%   82.87%   -0.40%     
==========================================
  Files          31       31              
  Lines        3759     3801      +42     
==========================================
+ Hits         3130     3150      +20     
- Misses        629      651      +22     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[yebai]

One high-level suggestion is to review DynamicPPL.LogDensityFunction and explore whether it can be shared with JuliaBUGS models. This could help reduce code redundancy.

cc @penelopeysm

[gdalle]

Is it really non-breaking to make such a big change?

[shravanngoswamii]

The adtype parameter is optional (defaults to nothing), all existing code works unchanged, so this is backward compatible.

However, if you prefer 0.11.0, I can change it.

[sunxd3]

I like this a lot, thanks! (The code really look fantastic and comprehensive.)

I had a similar PR locally, but didn't push through because I feels like we should first test against the AD backends so that when we say we support, we actually support.
There are some complications wrt AD backend too, for instance, only Enzyme and Mooncake support the sourcegen evaluation functions because the functions mutate. Mooncake and Enzyme probably work with the graph-based evaluation function, but the graph eval functions are only tested against ReverseDiff and ForwardDiff.

Another thing to consider is whether we want to drop the support for LogDensityProblemsAD (I think the current PR will, so as it is, it should be breaking).

@shravanngoswamii what do you think?

[yebai]

Another thing to consider is whether we want to drop the support for LogDensityProblemsAD (I think the current PR will, so as it is, it should be breaking).

Let's drop it unless there is a good reason.

[github-actions]

Benchmark results on macOS (aarch64)

BridgeStan not found at location specified by $BRIDGESTAN environment variable, downloading version 2.7.0 to /Users/runner/.bridgestan/bridgestan-2.7.0
Done!

Stan results:

Model Symbol	Parameters Int64	Density Time (µs) Float64	Density+Gradient Time (µs) Float64
rats	65	5.83333	8.083
pumps	12	0.947233	1.08955
bones	13	81.542	93.333
oxford	244	10.6045	12.667
epil	303	37.041	41.292
lsat	1006	106.042	147.792
schools	133	244.208	379.791
beetles	2	0.793919	0.905903
air	5	0.925	1.03817

JuliaBUGS Mooncake results:

Model Symbol	Parameters Int64	Density Time (µs) Float64	Density+Gradient Time (µs) Float64
rats	65	1.25543	7.95833
pumps	12	0.493274	1.97913
bones	33	53.291	280.25
oxford	244	23.542	60.9165
epil	303	7.861	36.958
lsat	1006	54.917	247.292
schools	133	123.666	1344.0
beetles	2	1.17954	5.1666
air	5	0.3125	1.21179

---

Benchmark results on Ubuntu (x64)

BridgeStan not found at location specified by $BRIDGESTAN environment variable, downloading version 2.7.0 to /home/runner/.bridgestan/bridgestan-2.7.0
Done!

Stan results:

Model Symbol	Parameters Int64	Density Time (µs) Float64	Density+Gradient Time (µs) Float64
rats	65	5.3682	7.63767
pumps	12	1.03443	1.21896
bones	13	69.922	86.383
oxford	244	14.142	16.581
epil	303	29.195	35.497
lsat	1006	160.372	197.422
schools	133	508.237	748.279
beetles	2	0.862212	0.979833
air	5	0.641413	0.794459

JuliaBUGS Mooncake results:

Model Symbol	Parameters Int64	Density Time (µs) Float64	Density+Gradient Time (µs) Float64
rats	65	1.92427	14.097
pumps	12	0.607604	2.42091
bones	33	47.709	371.609
oxford	244	24.195	75.689
epil	303	10.9855	58.179
lsat	1006	89.879	420.852
schools	133	101.912	1088.19
beetles	2	1.10515	4.71717
air	5	0.34425	1.29014

[yebai]

Anything missing here?

[shravanngoswamii]

Anything missing here?

Yes, We haven't tested against all mentioned AD Backends, somehow I was getting few errors, will have to look into it properly.

I had a similar PR locally, but didn't push through because I feels like we should first test against the AD backends so that when we say we support, we actually support.

Agree to you, we should test it before merging this!

Another thing to consider is whether we want to drop the support for LogDensityProblemsAD (I think the current PR will, so as it is, it should be breaking).

@sunxd3 I think we should drop the support for this, but it's still used in many parts of codebase. I don't have any strong opinions, still haven't explored JuliaBUGS in depth!