Integrate DifferentiationInterface.jl for gradient computation #397

shravanngoswamii · 2025-10-02T18:18:25Z

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 · 2025-10-02T18:21:55Z

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

shravanngoswamii · 2025-10-02T18:28:17Z

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

coveralls · 2025-10-02T18:46:22Z

Pull Request Test Coverage Report for Build 19888070954

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 83.484%

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 19887036408:	-0.4%
Covered Lines:	3321
Relevant Lines:	3978

💛 - Coveralls

codecov · 2025-10-02T18:51:27Z

Codecov Report

❌ Patch coverage is 55.55556% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.48%. Comparing base (acb8bf5) to head (a182845).

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.86%   83.48%   -0.39%     
==========================================
  Files          31       31              
  Lines        3936     3978      +42     
==========================================
+ Hits         3301     3321      +20     
- Misses        635      657      +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 · 2025-10-03T13:31:40Z

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

JuliaBUGS/docs/src/example.md

gdalle · 2025-10-06T07:22:05Z

JuliaBUGS/Project.toml

 name = "JuliaBUGS"
 uuid = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
-version = "0.10.3"
+version = "0.10.4"


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

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.

JuliaBUGS/docs/src/example.md

JuliaBUGS/src/JuliaBUGS.jl

sunxd3 · 2025-10-06T12:12:22Z

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 · 2025-10-06T12:32:24Z

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.

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

yebai · 2025-10-26T23:43:51Z

Anything missing here?

shravanngoswamii · 2025-10-27T02:26:48Z

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!

Julia 1.12 makes some changes to how world age is handled. JuliaBUGS use `eval` to generate node functions and this is problematic under 1.12. If a user uses JuliaBUGS `compile` in REPL, there is no problem, like before. This issue is particularly pronounced when using JuliaBUGS with a pattern like `compile -> LogDensityProblems.logdensity` within the same script. It can be pretty involved fixing the problem fundamentally. For now, this PR patches some of the tests to use `invokelatest` so they run. Other than it's nontrivial to fix the problem, we are likely seeing a change to the AD interface for JuliaBUGS following #397.

yebai · 2025-11-30T21:59:41Z

It's okay test ForwardDiff, and Mooncake only. Turing.jl will move towards Mooncake in the longer term.

Resolved conflicts: - Project.toml: Use main's version (0.11.1) - benchmark/benchmark.jl: Added comparison table functions from main - benchmark/run_benchmarks.jl: Use main's new benchmark output format - src/JuliaBUGS.jl: Merged main's evaluation mode docs with HEAD's adtype feature - test/ext/JuliaBUGS*Ext.jl: Keep HEAD's adtype compile API

BREAKING: Use compile(...; adtype=...) or BUGSModelWithGradient(model, adtype) instead of ADgradient()

sunxd3 · 2025-12-04T09:58:20Z

@shravanngoswamii I updated some code, can you give this a final review?

github-actions · 2025-12-04T10:10:41Z

Benchmark Results

Julia 1.11.7 on AMD EPYC 7763 64-Core Processor (Linux)

Ratio = JuliaBUGS/Stan (lower is better for JuliaBUGS)

Model	Stan Params	JBUGS Params	LD Ratio	Grad Ratio
rats	65	65	0.35x	1.69x
pumps	12	12	0.60x	1.81x
bones	13	33	0.67x	4.47x
oxford	244	244	1.65x	4.26x
epil	303	303	0.52x	1.59x
lsat	1006	1006	0.62x	1.89x
schools	133	133	0.23x	1.43x
beetles	2	2	1.36x	4.99x
air	5	5	0.52x	1.46x

Note: Stan benchmarks use hand-optimized Stan models, not direct BUGS translations. Comparison is illustrative only.

shravanngoswamii · 2025-12-04T12:07:36Z

@shravanngoswamii I updated some code, can you give this a final review?

@sunxd3 Thanks for taking this over. Everything looks good to me. Feel free to merge!

sunxd3 · 2025-12-04T12:23:43Z

Thanks!

shravanngoswamii added 2 commits October 2, 2025 22:38

Simplify the workflow for computing model gradients

8dfc800

update docs and bump version to 0.10.4

d77bbac

github-actions bot assigned shravanngoswamii Oct 2, 2025

run JuliaFormatter

199ba71

shravanngoswamii marked this pull request as ready for review October 2, 2025 18:28

shravanngoswamii requested review from sunxd3 and yebai and removed request for sunxd3 October 2, 2025 18:28

shravanngoswamii added the enhancement New feature or request label Oct 2, 2025

shravanngoswamii added 2 commits October 3, 2025 00:26

try to fix benchmark failures

1e515d1

format

f0135a8

shravanngoswamii force-pushed the sg/#380 branch from 511c6a6 to f0135a8 Compare October 2, 2025 19:15

gdalle reviewed Oct 6, 2025

View reviewed changes

add Mooncake and update docs

c80dc9f

penelopeysm reviewed Oct 6, 2025

View reviewed changes

JuliaBUGS/src/JuliaBUGS.jl Outdated Show resolved Hide resolved

shravanngoswamii and others added 2 commits October 7, 2025 12:13

remove symbols

ef36d7d

Update JuliaBUGS/test/ext/JuliaBUGSAdvancedHMCExt.jl

2f8b72d

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Merge branch 'main' into sg/#380

653540f

sunxd3 mentioned this pull request Nov 30, 2025

Use invokelatest in tests (because of Julia 1.12) #404

Merged

sunxd3 added 11 commits December 3, 2025 08:56

Refactor BUGSModelWithGradient and remove LogDensityProblemsAD support

a44fda5

BREAKING: Use compile(...; adtype=...) or BUGSModelWithGradient(model, adtype) instead of ADgradient()

Merge main into sg/#380

c130d78

Fix AD gradient function signature and add AD compatibility tests

7d033d8

format

94fb6c5

remove redundant lines in benchmark code

8aee4d2

various cleaning up

cd1df5e

Update examples to use new BUGSModelWithGradient API

a501476

update example doc

e059ad8

fix hmc test error

4509463

disable progress bar example doc

d6fc462

sunxd3 merged commit b40f54b into main Dec 4, 2025
27 of 28 checks passed

sunxd3 deleted the sg/#380 branch December 4, 2025 12:23

Integrate DifferentiationInterface.jl for gradient computation #397

Integrate DifferentiationInterface.jl for gradient computation #397

Uh oh!

Conversation

shravanngoswamii commented Oct 2, 2025

Uh oh!

github-actions bot commented Oct 2, 2025

Uh oh!

shravanngoswamii commented Oct 2, 2025

Uh oh!

coveralls commented Oct 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 19888070954

Details

💛 - Coveralls

Uh oh!

codecov bot commented Oct 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

yebai commented Oct 3, 2025

Uh oh!

Uh oh!

gdalle Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

shravanngoswamii Oct 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sunxd3 commented Oct 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yebai commented Oct 6, 2025

Uh oh!

yebai commented Oct 26, 2025

Uh oh!

shravanngoswamii commented Oct 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yebai commented Nov 30, 2025

Uh oh!

sunxd3 commented Dec 4, 2025

Uh oh!

github-actions bot commented Dec 4, 2025

Benchmark Results

Uh oh!

shravanngoswamii commented Dec 4, 2025

Uh oh!

sunxd3 commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

coveralls commented Oct 2, 2025 •

edited

Loading

codecov bot commented Oct 2, 2025 •

edited

Loading

shravanngoswamii Oct 6, 2025 •

edited

Loading

sunxd3 commented Oct 6, 2025 •

edited

Loading

shravanngoswamii commented Oct 27, 2025 •

edited

Loading

sunxd3 commented Dec 4, 2025 •

edited

Loading