Fix for #596

[torfjelde]

See #596 for a description of the issue. Fix #596

This PR automates the approach suggested by @willtebbutt, making it so that

@model demo(::Type{TV}=Vector{Float64}) where {TV}

is effectively converted into the current

@model demo(::DynamicPPL.TypeWrap{TV}=DynamicPPL.TypeWrap{Vector{Float64}}()) where {TV}

ensuring that model.args does not involve DataType, leading to significant improvements in type-stability for both the evaluator and the constructor.`

EDIT: Note that this change is a bit annoying as it does introduce "more magic" in DynamicPPL, but I'm personally happy with this for now as it technically fixes a bunch of type stability bugs. If we want to be explicit about what's happening, we can always just remove the magic at a later stage and tell the user to use TypeWrap instead of Type going forward.

[torfjelde]

The type-instabilities had been observed before too, e.g. see the test case now uncommented in this PR, but hadn't investigated it yet and thought it was for different reasons, so big thanks to @willtebbutt for discovering this!

[coveralls]

Pull Request Test Coverage Report for Build 8989979287

Details

• 16 of 16 (100.0%) changed or added relevant lines in 1 file are covered.
• 130 unchanged lines in 14 files lost coverage.
• Overall coverage decreased (-2.8%) to 78.739%

---

Files with Coverage Reduction	New Missed Lines	%
src/model.jl	1	89.22%
src/sampler.jl	1	94.12%
ext/DynamicPPLForwardDiffExt.jl	1	77.78%
src/contexts.jl	2	77.27%
src/logdensityfunction.jl	2	55.56%
src/compiler.jl	2	93.42%
src/prob_macro.jl	4	84.67%
src/abstract_varinfo.jl	5	82.68%
src/simple_varinfo.jl	6	82.61%
src/context_implementations.jl	8	62.23%

Totals
Change from base Build 8986295028:	-2.8%
Covered Lines:	2785
Relevant Lines:	3537

---

💛 - Coveralls

[torfjelde]

@yebai @devmotion @sunxd3 thoughts on this?

I think the main issue is: do we want to introduce more magic (here we perform a simple transformation on very specific parameters)? As outlined above, I'm pro doing it for now as it will have significant impact on user code (type-stability and perf wise + will make some AD backends much faster) and I consider this a bugfix (given how we spout this pattern as a way of achieving signfiicnat performance gains). Long-term we might want to make these things explicit, but for now I like this.

[devmotion]

Great!

we can always just remove the magic at a later stage and tell the user to use TypeWrap instead of Type going forward.

I think ideally we would not have to use TypeWrap internally. But IMO the "magic" is better than exposing TypeWrap, I think we should not require users to know about such details (and even if it would be in the docs I am sure the majority of users would unknowingly use Type).

[willtebbutt]

Does this just require a patch bump?

[torfjelde]

Cheers @devmotion !

Does this just require a patch bump?

Done.

[torfjelde]

I've been running some benchmarks locally with ForwardDiff, ReverseDiff, and Zygote, and it doesn't seem like there are any perf improvements 🙃 With the exception of models using @submodel, which was expected to improve signfiicantly (since in that case the return-value being inferred is mucho bueno).

EDIT: But there consistently less memory usage, so it's at least doing something:)

[sunxd3]

Love this!

[torfjelde]

So integration tests are failing because we're explicitly testing passing in types as arguments. It does raise the question as to whether we should make this breaking or not 🤔

Thoughts?

[devmotion]

Can we have both? That is, could we forward f(arg1, ::Type{TV}, ...) to f(arg1, ::TypeWrap{TV}, ...)?

[torfjelde]

Can we have both? That is, could we forward f(arg1, ::Type{TV}, ...) to f(arg1, ::TypeWrap{TV}, ...)?

Was thinking about the same. But so now we generate 3 methods instead? And should we issue a dep warning?

Or do we also add an evaluator with this, leading to 4 methods?

[torfjelde]

Hmm on a second thought, I don't think we can 😕

If we try to keep the old method around, we will define the default method twice. Only if we remove the default values for the Type version could we do something like that, but then we're already probably breaking stuff so uncertain if that's worth it 😕

[torfjelde]

tbf there's probably no one using this functionality outside of Turing.jl's own tests, so probably isn't a huge problem to make a breaking change to it

[yebai]

tbf there's probably no one using this functionality outside of Turing.jl's own tests, so probably isn't a huge problem to make a breaking change to it

Sounds good.

[yebai]

Is anything missing here? If not, let's try to get this merged.

[devmotion]

I think this should probably be applied to keyword arguments as well?

Suggested change

	kwargs = modeldef[:kwargs]
	kwargs = transform_args(modeldef[:kwargs])

(not sure if it can be applied directly here in this way)

[torfjelde]

Why? Do we care if we specialize on the kwargs?

[devmotion]

IIRC we store them in the same way as pos arguments in the model.

[devmotion]

So the Model will generally still operate with DataTypes unless we apply the same fix to kwarga.

[torfjelde]

No you're right 👍 Working on it now

[torfjelde]

Added this 👍

[torfjelde]

Only question: do we make this a breaking release or not? @devmotion thoughts?

Technically it's breaking so I am leaning towards this.

[yebai]

Probably better to make this breaking release since it breaks Turing tests anyway

[devmotion]

I'd also make it a breaking release.

[yebai]

It seems auto merging is quite fragile — it needs disabling and then re-enabling to fix unknown failures.

[torfjelde]

Please let the author of the PR perform the merge @yebai ; in this case it was fine, but it's not always the case.