Hello, congrats on the new package! 😄
I wanted to ask a couple of things about
|
# Build the LogDensityProblems-compatible gradient object |
|
ld = DynamicPPL.LogDensityFunction(turing_model) |
|
ldg = LogDensityProblemsAD.ADgradient(ad_backend, ld) |
|
|
|
caps = LogDensityProblems.capabilities(ldg) |
|
caps isa LogDensityProblems.LogDensityOrder{0} && error( |
|
"AD gradient setup failed. The wrapped model must support gradients. " * |
|
"Ensure your ad_backend is compatible.", |
|
) |
Firstly, when constructing a LogDensityFunction, to get AD you don't need to wrap it in ADgradient (in fact we used to do that and then stopped doing that a while ago). You can directly pass the adtype kwarg, see e.g. https://turinglang.org/docs/usage/vectorisation/#automatic-differentiation and that will directly give you something that you can run logdensity_and_gradient on.
Turing/DynamicPPL does a fair bit of testing with Enzyme, but all the tests go via this adtype kwarg path (in effect, via DifferentiationInterface) and we don't do any testing with LogDensityProblemsAD.jl (which has custom Enzyme bindings and therefore circumvents DI). LDPAD certainly works, but I reckon you might get some extra security from the existing tests if you used the DI path :-).
Secondly, do you require or prefer that the model is operating in unconstrained space? Just calling LogDensityFunction(model) will run in the 'original' space, for example, so I wonder if maybe as a default you might want to avoid that
julia> using DynamicPPL, Distributions, LogDensityProblems
julia> @model f() = x ~ Beta(2, 2); ldf = LogDensityFunction(f());
julia> LogDensityProblems.logdensity(ldf, [-0.4])
-Inf
julia> ldf_uncons = LogDensityFunction(f(), getlogjoint_internal, LinkAll());
julia> LogDensityProblems.logdensity(ldf_uncons, [-0.4])
-1.0603015403717557
Hello, congrats on the new package! 😄
I wanted to ask a couple of things about
ParallelMCMC.jl/ext/DynamicPPLExt.jl
Lines 48 to 56 in 882b6b4
Firstly, when constructing a
LogDensityFunction, to get AD you don't need to wrap it inADgradient(in fact we used to do that and then stopped doing that a while ago). You can directly pass theadtypekwarg, see e.g. https://turinglang.org/docs/usage/vectorisation/#automatic-differentiation and that will directly give you something that you can runlogdensity_and_gradienton.Turing/DynamicPPL does a fair bit of testing with Enzyme, but all the tests go via this
adtypekwarg path (in effect, via DifferentiationInterface) and we don't do any testing with LogDensityProblemsAD.jl (which has custom Enzyme bindings and therefore circumvents DI). LDPAD certainly works, but I reckon you might get some extra security from the existing tests if you used the DI path :-).Secondly, do you require or prefer that the model is operating in unconstrained space? Just calling
LogDensityFunction(model)will run in the 'original' space, for example, so I wonder if maybe as a default you might want to avoid that