Simplify the workflow for computing model gradients #380
Description
Current JuliaBUGS workflow for gradient-based MCMC:
JuliaBUGS.jl/JuliaBUGS/test/ext/JuliaBUGSAdvancedHMCExt.jl
Lines 3 to 13 in 533b93b
It can be simplified into
model_def = @bugs begin
x[1:2] ~ dmnorm(mu[:], sigma[:, :])
x[3] ~ dnorm(0, 1)
y = x[1] + x[3]
end
data = (mu=[0, 0], sigma=[1 0; 0 1])
- model = compile(model_def, data)
- ad_model = ADgradient(AutoReverseDiff(compile=true), model)
- D = LogDensityProblems.dimension(model)
+ model = compile(model_def, data) # `adtype` defaults to `nothing`
+ ad_model = compile(model_def, data; adtype=AutoReverseDiff(compile=true))
samples_and_stats = AbstractMCMC.sample(
StableRNG(1234),
ad_model,
NUTS(0.8),
2000;
progress=false,
chain_type=Chains,
n_adapts=1000,
init_params=rand(D),
discard_initial=n_adapts,
)where compile(model_def, data; adtype=AutoReverseDiff()) returns an ADgradient object directly when adtype keyword argument is passed (adtype defaults to nothing). In addition, the LogDensityProblems.dimension function should be automatically defined under the hood, too.