Using JET.jl to determine if typed varinfo is okay

[torfjelde]

After a quick experiment with JET.jl I found some bugs in DynamicPPL.jl (#726), but also realized that we can JET.jl to properly check whether the a given model supports the usage of TypedVarInfo rather than requiring UntypedVarInfo.

This has been a looooooong standing issue, and this seems to work really, really well.

The problem

In Turing.jl, we use TypedVarInfo almost everywhere due to the performance charactersitics that come with it. The problem is that we do so by simply evaluating the given model once and then using the resulting (hopefully, concretety typed) varinfo for all subsequent computations. This works nicely for most typical models, but fails horribly (and uninformably) for a good chunk of models, such as

@model function demo1()
    x ~ Bernoulli()
    if x
        y ~ Normal()
    else
        z ~ Normal()
    end
end

Here we will execute the model once and get, say, a TypedVarInfo containing the variables x and y (because x happend to result in a true sample). If we then re-use this varinfo for sampling, we will ofc run into issues since z is nowhere to be seen.

Technically we can handle this by just widing the container a bit, but if we do that, we need to cpature the new varinfo, which isn't always possible, e.g. when using the LogDensityFunction in a sampler.

As a result, we have a lot of code that just makes the assumption "surely this model is 'static' in what variables and types it contains", which can sometimes be false.

The solution

This PR introduces a determine_varinfo method, which can automagically figure out whether we can use the type stable varinfo properly (i.e. without having to always capture the resulting varinfo, etc.) or if we need to use the untyped varinfo using abstract interpretation offered by JET.jl, all done statically.

Effectively what determine_varinfo does is:

1. Execute the model once with to get the typed varinfo.
2. Using JET.jl, statically check if we can run into type issues, e.g. container of NamedTuple{(:x, :y)} cannot handle the value for z being updated (because the entry does not exist).
3. If we do run into errors, we return an untyped varinfo. If we don't, we return a typed one.

Note that this method doesn't say anything about whether there might be type instabilities; this only checks if we would encounter errors. We can also use JET to check type instabilites, etc., but I think that's a separate functionality and thus PR.

[torfjelde]

See the tests for what we can properly check here. It honestly seems really good for our purposes 👀

[yebai]

That seems like an elegant trick!

[codecov]

Codecov Report

Attention: Patch coverage is 89.28571% with 3 lines in your changes missing coverage. Please review.

Project coverage is 86.49%. Comparing base (0548ddf) to head (4a17e82).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
src/experimental.jl	60.00%	2 Missing ⚠️
src/DynamicPPL.jl	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #728      +/-   ##
==========================================
+ Coverage   86.34%   86.49%   +0.15%     
==========================================
  Files          34       36       +2     
  Lines        4254     4272      +18     
==========================================
+ Hits         3673     3695      +22     
+ Misses        581      577       -4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[coveralls]

Pull Request Test Coverage Report for Build 12236332369

Details

• 25 of 28 (89.29%) changed or added relevant lines in 4 files are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage increased (+0.2%) to 86.493%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/DynamicPPL.jl	4	5	80.0%
src/experimental.jl	3	5	60.0%

Totals
Change from base Build 12224895333:	0.2%
Covered Lines:	3695
Relevant Lines:	4272

---

💛 - Coveralls

[torfjelde]

This honestly seem to work really well. I've now made it so that default_sampler and LogDensityFunction also makes use of this. The question is just how well it works with Turing.jl (will try this now).

[torfjelde]

Added some docs 👍

[mhauru]

Thanks @torfjelde, this seems really handy. Had a few localised comments, nothing major.

[mhauru]

Could there be some way to test this? I do see that it's tricky. I'm a bit uncomfortable having this in without any testing.

[torfjelde]

Yeah was thinking the same. We could put in a test strictly before loading JET.jl ofc. It's a bit messy, but seems like the best way 😕

[mhauru]

Does wrapping the tests in separate modules save us?

[torfjelde]

Nah. AFAIK extensions trigger if the package is loaded at any point, e.g. even if a dep loads it

[torfjelde]

It's also a thing where it doesn't seem like we can nicely get the resulting error message (the error hint is not in the msg of the error or something). So I think we just leave this for now 😕

[mhauru]

Sure, it does seem nasty to test for. Have you tried locally that it does what you expect?

[torfjelde]

This is the example mentioned above @mhauru :)

[torfjelde]

You happy with this now @mhauru ?:) It's only failiing because of the x86 OOM thingy

[mhauru]

Yup, am happy, thanks!