FlexiChains by default #44
There are no files selected for viewing
|
penelopeysm marked this conversation as resolved.
|
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,7 +4,7 @@ using ParallelMCMC | |
| using ADTypes: ADTypes | ||
| using DynamicPPL: DynamicPPL | ||
| using AbstractMCMC: AbstractMCMC | ||
| using FlexiChains: FlexiChain, VarName, VNChain, SymChain | ||
| using LogDensityProblems: LogDensityProblems | ||
|
|
||
| """ | ||
|
|
@@ -14,13 +14,12 @@ Convenience constructor: wraps a DynamicPPL/Turing `@model` directly as a | |
| `DensityModel`, automatically extracting parameter names and wiring up gradient | ||
| computation via DynamicPPL's `adtype` interface. | ||
|
|
||
| Requires `DynamicPPL` and `LogDensityProblems` to be loaded (these are the weak-dependency | ||
| triggers for this extension), plus any AD backend that is used. | ||
|
|
||
| # Example | ||
| ```julia | ||
| using Turing, ParallelMCMC, FlexiChains | ||
|
penelopeysm marked this conversation as resolved.
|
||
|
|
||
| @model function mymodel(y) | ||
| μ ~ Normal(0, 1) | ||
|
|
@@ -31,7 +30,7 @@ end | |
| # and `using` the corresponding package (Enzyme, Mooncake). | ||
| model = DensityModel(mymodel(1.5)) | ||
| chain = sample(model, AdaptiveMALASampler(0.3; n_warmup=500), 2_000; | ||
| chain_type=FlexiChains.VNChain, discard_warmup=true, progress=true) | ||
| ``` | ||
| """ | ||
| function ParallelMCMC.DensityModel( | ||
|
Comment on lines
35
to
36
Owner
There was a problem hiding this comment. So right now I have ForwardDiff as the dfault for this ext because of issue #25 and ForwardDiff worked as a workaround. I failed to document this so thats my b. though I'm not sure it makes total sense to have this as a default just given there is other backends i didn't try. So at the very least we need to either re-add in the ForwardDiff as a gate for this ext or just remove this
Contributor
Author
There was a problem hiding this comment. No worries. So personally I don't think it's necessary to have ForwardDiff as an extension trigger -- for one, maybe people don't actually want to use ForwardDiff (in which case forcing them to load it is one extra dep), and secondly it reduces discoverability. After all the main package is also agnostic and doesn't load any AD backend as a hard dep so I think the extension can do the same thing.
Owner
There was a problem hiding this comment. yes i agree let's remove as hard dep and remove the default
Contributor
Author
There was a problem hiding this comment. Hmm I was thinking it's okay to keep AutoForwardDiff as the default (it would be a bit annoying as if you don't import ForwardDiff and run it it would error, but that's no different from the main package). However if I'm not mistaken, you are suggesting to leave backend unspecified in the DynamicPPL extension, and let the user pass it via the sampler instead? If that's the case I do actually think that that's better but maybe we can do it in a separate PR?
Owner
There was a problem hiding this comment. yes that was my thinking to prevent the error, but yes thats fine we can push till next PR
Contributor
Author
There was a problem hiding this comment. Okay so for now I just reverted this so that ForwardDiff is still an extension trigger (restoring the old behaviour) and we could look at it later. I haven't thought too much about it but it might be a bit of a faff as you have to re-wrap the DynamicPPL model with the adtype from the sampler hence why I thought it made sense to defer. |
||
|
|
@@ -106,40 +105,47 @@ for (Ttrans, Tspl, Tstate) in ( | |
| model::DensityModelLDF, | ||
| spl::$Tspl, | ||
| state::$Tstate, | ||
| chain_type::Type{VNChain}; | ||
| discard_warmup::Bool=false, | ||
|
penelopeysm marked this conversation as resolved.
|
||
| kwargs..., | ||
| ) | ||
| ts = discard_warmup ? filter(t -> !is_warmup(t), ts) : ts | ||
| pwss = map(ts) do t | ||
| # Note: This assumes that there is always a field called t.x. This is currently true | ||
| # of all samplers in ParallelMCMC | ||
| DynamicPPL.ParamsWithStats(t.x, model.logdensity.ld, getstats(t)) | ||
| end | ||
| return AbstractMCMC.from_samples(VNChain, hcat(pwss)) | ||
| end | ||
| end | ||
|
|
||
| @eval begin | ||
| function AbstractMCMC.bundle_samples( | ||
| ts::Vector{<:$Ttrans}, | ||
| model::DensityModelLDF, | ||
| spl::$Tspl, | ||
| state::$Tstate, | ||
| chain_type::Type{SymChain}; | ||
| kwargs..., | ||
| ) | ||
| throw(ArgumentError("FlexiChains.SymChain is not supported for DynamicPPL models; please use VNChain instead.")) | ||
|
Comment on lines
+123
to
+131
Contributor
Author
There was a problem hiding this comment. So I was going to add a test for this but before I do that, I realised that this might be kind of annoying behaviour, in that if you specify the wrong chain type it will do all the sampling but only error at the end. That's no worse than the current behaviour ofc (current behaviour is to do all the sampling and then have an inscrutable MethodError at the end). Just wanted to check you're okay with that first!
Owner
There was a problem hiding this comment. hmm yeah a bit annoying, but as you said better than the earlier behavior with a legible error now. I say this is good w/ me and can be revisited later if necessary. Go ahead and add some tests! |
||
| end | ||
| end | ||
| end | ||
|
|
||
| function ParallelMCMC._construct_flexichain( | ||
| ::Type{VarName}, | ||
| vals::AbstractMatrix{<:Real}, | ||
| internals::AbstractMatrix{<:Real}, | ||
| ::Any, | ||
| internal_names::Vector{Symbol}, | ||
| model::DensityModelLDF, | ||
| ) | ||
| pwss = map(zip(eachrow(vals), eachrow(internals))) do (val, internal) | ||
| stats = NamedTuple{Tuple(internal_names)}(internal) | ||
| DynamicPPL.ParamsWithStats(val, model.logdensity.ld, stats) | ||
|
penelopeysm marked this conversation as resolved.
|
||
| end | ||
| return AbstractMCMC.from_samples(VNChain, hcat(pwss)) | ||
| end | ||
|
|
||
| end # module | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should still declare a compat v for this even if removed from the Ext
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh, is it known to not work with ForwardDiff v0?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
hmm i guess not known, but i haven't tested so maybe its fine