Simplify the workflow for computing model gradients

Author: shravanngoswamii

.gitignore

@@ -29,4 +29,6 @@ Manifest.toml
 *.local.*
 
 # gitingest generated files
-digest.txt
+digest.txt
+
+tmp/

JuliaBUGS/Project.toml

@@ -9,6 +9,7 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -52,6 +53,7 @@ AdvancedHMC = "0.6, 0.7, 0.8"
 AdvancedMH = "0.8"
 BangBang = "0.4.1"
 Bijectors = "0.13, 0.14, 0.15.5"
+DifferentiationInterface = "0.7"
 Distributions = "0.23.8, 0.24, 0.25"
 Documenter = "0.27, 1"
 GLMakie = "0.10, 0.11, 0.12, 0.13"

JuliaBUGS/ext/JuliaBUGSAdvancedHMCExt.jl

@@ -3,10 +3,12 @@ module JuliaBUGSAdvancedHMCExt
 using AbstractMCMC
 using AdvancedHMC
 using ADTypes
+import DifferentiationInterface as DI
 using JuliaBUGS
-using JuliaBUGS: BUGSModel, getparams, initialize!
+using JuliaBUGS: BUGSModel, BUGSModelWithGradient, getparams, initialize!
 using JuliaBUGS.LogDensityProblems
 using JuliaBUGS.LogDensityProblemsAD
+using JuliaBUGS.Model: _logdensity_switched
 using JuliaBUGS.Random
 using MCMCChains: Chains
 
@@ -40,10 +42,13 @@ end
 function _gibbs_internal_hmc(
     rng::Random.AbstractRNG, cond_model::BUGSModel, sampler, ad_backend, state
 )
-    # Wrap model with AD gradient computation
-    logdensitymodel = AbstractMCMC.LogDensityModel(
-        LogDensityProblemsAD.ADgradient(ad_backend, cond_model)
+    # Create gradient model on-the-fly using DifferentiationInterface
+    x = getparams(cond_model)
+    prep = DI.prepare_gradient(
+        _logdensity_switched, ad_backend, x, DI.Constant(cond_model)
     )
+    ad_model = BUGSModelWithGradient(ad_backend, prep, cond_model)
+    logdensitymodel = AbstractMCMC.LogDensityModel(ad_model)
 
     # Take HMC/NUTS step
     if isnothing(state)
@@ -53,7 +58,7 @@ function _gibbs_internal_hmc(
             logdensitymodel,
             sampler;
             n_adapts=0,  # Disable adaptation within Gibbs
-            initial_params=getparams(cond_model),
+            initial_params=x,
         )
     else
         # Use existing state for subsequent steps
@@ -67,7 +72,7 @@ end
 
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:AdvancedHMC.Transition},
-    logdensitymodel::AbstractMCMC.LogDensityModel{<:LogDensityProblemsAD.ADGradientWrapper},
+    logdensitymodel::AbstractMCMC.LogDensityModel{<:BUGSModelWithGradient},
     sampler::AdvancedHMC.AbstractHMCSampler,
     state,
     chain_type::Type{Chains};
@@ -98,4 +103,35 @@ function AbstractMCMC.bundle_samples(
     )
 end
 
+# Keep backward compatibility with LogDensityProblemsAD wrapper
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:AdvancedHMC.Transition},
+    logdensitymodel::AbstractMCMC.LogDensityModel{<:LogDensityProblemsAD.ADGradientWrapper},
+    sampler::AdvancedHMC.AbstractHMCSampler,
+    state,
+    chain_type::Type{Chains};
+    discard_initial=0,
+    thinning=1,
+    kwargs...,
+)
+    param_samples = [t.z.θ for t in ts]
+
+    stats_names = collect(keys(merge((; lp=ts[1].z.ℓπ.value), AdvancedHMC.stat(ts[1]))))
+    stats_values = [
+        vcat(ts[i].z.ℓπ.value, collect(values(AdvancedHMC.stat(ts[i])))) for
+        i in eachindex(ts)
+    ]
+
+    # Delegate to gen_chains for proper parameter naming from BUGSModel
+    return JuliaBUGS.gen_chains(
+        logdensitymodel,
+        param_samples,
+        stats_names,
+        stats_values;
+        discard_initial=discard_initial,
+        thinning=thinning,
+        kwargs...,
+    )
+end
+
 end

JuliaBUGS/ext/JuliaBUGSAdvancedMHExt.jl

@@ -3,10 +3,12 @@ module JuliaBUGSAdvancedMHExt
 using AbstractMCMC
 using AdvancedMH
 using ADTypes
+import DifferentiationInterface as DI
 using JuliaBUGS
-using JuliaBUGS: BUGSModel, getparams, initialize!
+using JuliaBUGS: BUGSModel, BUGSModelWithGradient, getparams, initialize!
 using JuliaBUGS.LogDensityProblems
 using JuliaBUGS.LogDensityProblemsAD
+using JuliaBUGS.Model: _logdensity_switched
 using JuliaBUGS.Random
 using MCMCChains: Chains
 
@@ -52,10 +54,13 @@ end
 function _gibbs_internal_mh(
     rng::Random.AbstractRNG, cond_model::BUGSModel, sampler, ad_backend, state
 )
-    # Wrap model with AD gradient computation for gradient-based proposals
-    logdensitymodel = AbstractMCMC.LogDensityModel(
-        LogDensityProblemsAD.ADgradient(ad_backend, cond_model)
+    # Create gradient model on-the-fly using DifferentiationInterface
+    x = getparams(cond_model)
+    prep = DI.prepare_gradient(
+        _logdensity_switched, ad_backend, x, DI.Constant(cond_model)
     )
+    ad_model = BUGSModelWithGradient(ad_backend, prep, cond_model)
+    logdensitymodel = AbstractMCMC.LogDensityModel(ad_model)
 
     # Take MH step with gradient information
     if isnothing(state)
@@ -64,7 +69,7 @@ function _gibbs_internal_mh(
             logdensitymodel,
             sampler;
             n_adapts=0,  # Disable adaptation within Gibbs
-            initial_params=getparams(cond_model),
+            initial_params=x,
         )
     else
         t, s = AbstractMCMC.step(rng, logdensitymodel, sampler, state; n_adapts=0)
@@ -103,6 +108,32 @@ function AbstractMCMC.bundle_samples(
     )
 end
 
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:AdvancedMH.Transition},
+    logdensitymodel::AbstractMCMC.LogDensityModel{<:BUGSModelWithGradient},
+    sampler::AdvancedMH.MHSampler,
+    state,
+    chain_type::Type{Chains};
+    discard_initial=0,
+    thinning=1,
+    kwargs...,
+)
+    param_samples = [t.params for t in ts]
+    stats_names = [:lp]
+    stats_values = [[t.lp] for t in ts]
+
+    return JuliaBUGS.gen_chains(
+        logdensitymodel,
+        param_samples,
+        stats_names,
+        stats_values;
+        discard_initial=discard_initial,
+        thinning=thinning,
+        kwargs...,
+    )
+end
+
+# Keep backward compatibility with LogDensityProblemsAD wrapper
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:AdvancedMH.Transition},
     logdensitymodel::AbstractMCMC.LogDensityModel{<:LogDensityProblemsAD.ADGradientWrapper},

JuliaBUGS/ext/JuliaBUGSMCMCChainsExt.jl

@@ -2,7 +2,7 @@ module JuliaBUGSMCMCChainsExt
 
 using AbstractMCMC
 using JuliaBUGS
-using JuliaBUGS: BUGSModel, find_generated_quantities_variables, evaluate!!, getparams
+using JuliaBUGS: BUGSModel, BUGSModelWithGradient, find_generated_quantities_variables, evaluate!!, getparams
 using JuliaBUGS.AbstractPPL
 using JuliaBUGS.Accessors
 using JuliaBUGS.LogDensityProblemsAD
@@ -21,6 +21,20 @@ function JuliaBUGS.gen_chains(
     )
 end
 
+function JuliaBUGS.gen_chains(
+    model::AbstractMCMC.LogDensityModel{<:BUGSModelWithGradient},
+    samples,
+    stats_names,
+    stats_values;
+    kwargs...,
+)
+    # Extract BUGSModel from gradient wrapper
+    bugs_model = model.logdensity.base_model
+
+    return JuliaBUGS.gen_chains(bugs_model, samples, stats_names, stats_values; kwargs...)
+end
+
+# Keep backward compatibility with LogDensityProblemsAD wrapper
 function JuliaBUGS.gen_chains(
     model::AbstractMCMC.LogDensityModel{<:LogDensityProblemsAD.ADGradientWrapper},
     samples,

JuliaBUGS/src/JuliaBUGS.jl

@@ -6,6 +6,7 @@ using Accessors
 using ADTypes
 using BangBang
 using Bijectors: Bijectors
+using DifferentiationInterface
 using Distributions
 using Graphs, MetaGraphsNext
 using LinearAlgebra
@@ -17,6 +18,7 @@ using Serialization: Serialization
 using StaticArrays
 
 import Base: ==, hash, Symbol, size
+import DifferentiationInterface as DI
 import Distributions: truncated
 
 export @bugs
@@ -239,20 +241,56 @@ function validate_bugs_expression(expr, line_num)
 end
 
 """
-    compile(model_def, data[, initial_params]; skip_validation=false)
+    compile(model_def, data[, initial_params]; skip_validation=false, adtype=nothing)
 
 Compile the model with model definition and data. Optionally, initializations can be provided.
 If initializations are not provided, values will be sampled from the prior distributions.
 
 By default, validates that all functions in the model are in the BUGS allowlist (suitable for @bugs macro).
 Set `skip_validation=true` to skip validation (for @model macro usage).
+
+If `adtype` is provided, returns a `BUGSModelWithGradient` that supports gradient-based MCMC 
+samplers like HMC/NUTS. The gradient computation is prepared during compilation for optimal performance.
+
+# Arguments
+- `model_def::Expr`: Model definition from @bugs macro
+- `data::NamedTuple`: Observed data
+- `initial_params::NamedTuple=NamedTuple()`: Initial parameter values (optional)
+- `skip_validation::Bool=false`: Skip function validation (for @model macro)
+- `eval_module::Module=@__MODULE__`: Module for evaluation
+- `adtype`: AD backend specification. Can be:
+  - `AutoReverseDiff(compile=true)` - ReverseDiff with tape compilation (fastest)
+  - `AutoReverseDiff(compile=false)` - ReverseDiff without compilation
+  - `:ReverseDiff` - Shorthand for `AutoReverseDiff(compile=true)`
+  - `:ForwardDiff` - Shorthand for `AutoForwardDiff()`
+  - `:Zygote` - Shorthand for `AutoZygote()`
+  - Any other `ADTypes.AbstractADType`
+
+# Examples
+```julia
+# Basic compilation
+model = compile(model_def, data)
+
+# With gradient support using explicit ADType
+model = compile(model_def, data; adtype=AutoReverseDiff(compile=true))
+
+# With gradient support using symbol shorthand
+model = compile(model_def, data; adtype=:ReverseDiff)  # Same as above
+
+# Using ForwardDiff for small models
+model = compile(model_def, data; adtype=:ForwardDiff)
+
+# Sample with NUTS
+chain = AbstractMCMC.sample(model, NUTS(0.8), 1000)
+```
 """
 function compile(
     model_def::Expr,
     data::NamedTuple,
     initial_params::NamedTuple=NamedTuple();
     skip_validation::Bool=false,
     eval_module::Module=@__MODULE__,
+    adtype::Union{Nothing,ADTypes.AbstractADType,Symbol}=nothing,
 )
     # Validate functions by default (for @bugs macro usage)
     # Skip validation only for @model macro
@@ -281,7 +319,65 @@ function compile(
             values(eval_env),
         ),
     )
-    return BUGSModel(g, nonmissing_eval_env, model_def, data, initial_params)
+    base_model = BUGSModel(g, nonmissing_eval_env, model_def, data, initial_params)
+    
+    # If adtype provided, wrap with gradient capabilities
+    if adtype !== nothing
+        # Convert symbol to ADType if needed
+        adtype_obj = _resolve_adtype(adtype)
+        return _wrap_with_gradient(base_model, adtype_obj)
+    end
+    
+    return base_model
+end
+
+"""
+    _resolve_adtype(adtype) -> ADTypes.AbstractADType
+
+Convert symbol shortcuts to ADTypes, or return the ADType as-is.
+
+Supported symbol shortcuts:
+- `:ReverseDiff` -> `AutoReverseDiff(compile=true)`
+- `:ForwardDiff` -> `AutoForwardDiff()`
+- `:Zygote` -> `AutoZygote()`
+- `:Enzyme` -> `AutoEnzyme()`
+"""
+function _resolve_adtype(adtype::Symbol)
+    if adtype === :ReverseDiff
+        return ADTypes.AutoReverseDiff(compile=true)
+    elseif adtype === :ForwardDiff
+        return ADTypes.AutoForwardDiff()
+    elseif adtype === :Zygote
+        return ADTypes.AutoZygote()
+    elseif adtype === :Enzyme
+        return ADTypes.AutoEnzyme()
+    else
+        error("Unknown AD backend symbol: $adtype. " *
+              "Supported symbols: :ReverseDiff, :ForwardDiff, :Zygote, :Enzyme. " *
+              "Or use an ADTypes object like AutoReverseDiff(compile=true).")
+    end
+end
+
+# Pass through ADTypes objects unchanged
+_resolve_adtype(adtype::ADTypes.AbstractADType) = adtype
+
+# Helper function to prepare gradient - separated to handle world age issues
+function _wrap_with_gradient(base_model::Model.BUGSModel, adtype::ADTypes.AbstractADType)
+    # Get initial parameters for preparation
+    # Use invokelatest to handle world age issues with generated functions
+    x = Base.invokelatest(getparams, base_model)
+    
+    # Prepare gradient using DifferentiationInterface
+    # Use invokelatest to handle world age issues when calling logdensity during preparation
+    prep = Base.invokelatest(
+        DI.prepare_gradient,
+        Model._logdensity_switched,
+        adtype,
+        x,
+        DI.Constant(base_model)
+    )
+    
+    return Model.BUGSModelWithGradient(adtype, prep, base_model)
 end
 # function compile(
 #     model_str::String,

JuliaBUGS/src/model/Model.jl

@@ -2,8 +2,10 @@ module Model
 
 using Accessors
 using AbstractPPL
+using ADTypes
 using BangBang
 using Bijectors
+import DifferentiationInterface as DI
 using Distributions
 using Graphs
 using LinearAlgebra
@@ -21,5 +23,6 @@ include("logdensityproblems.jl")
 export parameters, variables, initialize!, getparams, settrans, set_evaluation_mode
 export regenerate_log_density_function, set_observed_values!
 export evaluate_with_rng!!, evaluate_with_env!!, evaluate_with_values!!
+export BUGSModelWithGradient, _logdensity_switched
 
 end # Model