[WIP] Taking stochastic control flow seriously

[mohamed82008]

This PR implements the data structures and methods needed to handle models with stochastic control flow of the form described in #25. To this end, I implemented the MixedVarInfo data structure which combines a TypedVarInfo and an UntypedVarInfo using the latter for symbols which have not been seen before. In the future, we can consider moving more of UntypedVarInfo to TypedVarInfo in an incremental sample function in Turing, where new type information is specialized upon for faster sampling. Currently, this PR is still a proof of concept. More things fail than work but once I implement the following, things should be much better. I moved the Turing src from Turing to the test folder and kept it in sync with Turing#mt/mixed_varinfo which I am using locally for tests.

For correctness and for things to work, the following still needs to be done:

1. Implement more methods for MixedVarInfo, namely _getvns and set_namedtuple!. Currently I just forward to the tvi::TypedVarInfo implementation which is strictly speaking not correct.
2. Implement tonamedtuple properly for UntypedVarInfo
3. Fix transition_save!. transition_save! currently errors when it sees named tuples with different number of fields.
4. Add many more tests:
   • MixedVarInfo behaviour tests for dynamic models with fixed symbols and dynamic models with a stochastic number of symbols.
   • Type stability tests
   • AD tests

For performance, the following needs to be done:

1. Model refactor
   • Store the LHS of ~ in a space field in the Model struct
   • Define getspace for Model to return the stored space field
   • Define getparams for model to return the symbols that are:
     1. In space and not in args, or
     2. In args but can admit the value missing
   • Define getobservations to return the symbols that are:
     1. In space,
     2. In args, and
     3. Cannot admit the value missing.
   • Define gethyperparameters to return the symbols that are:
     1. Not in space, and
     2. Are in args.
2. NamedDist refactor
   • Define a bijector, reconstruct and vectorize for NamedDist
   • Store the NamedDist in varinfo as is - don't replace the varname in assume as we do now
   • Use the varname in the NamedDist when converting the varinfo to a NamedTuple/Chain
3. AbstractSampler refactor
   • When constructing Sampler from alg and model::Model, where alg has no space defined, use getparams(model) as the space in Sampler.
   • Let SampleFromUniform and SampleFromPrior have a space type parameter that defaults to ().
   • Overload model(vi, ::Union{SampleFromPrior, SampleFromUniform}, ...) to replace the SampleFromPrior{()}() or SampleFromUniform{()}() with an empty space by SampleFromPrior{getparams(model)}() and SampleFromUniform{getparams(model)}() respectively.
   • Do the same in sample if sampling from SampleFromPrior.

That's a long list, so I expect this to be a huge PR. But I think it's worth it :) The tests here should pass except for the test in #104. This is not a sign that everything is correct though, I still need to add and probably fix many more tests.

[devmotion]

Regarding

Fix transition_save!. transition_save! currently errors when it sees named tuples with different number of fields.

We should just use BangBang in AbstractMCMC.

[devmotion]

BTW I think it is very difficult to review a PR and think through everything thoroughly if it affects so many files and implements so many changes at once. Is it possible to split this PR into smaller chunks that could be applied iteratively, focusing on one problem at a time? 😃

[mohamed82008]

Hmm to avoid any performance drops in master any one point, we will need to follow the above list in reverse. The Model and NamedDist refactors need to happen in the same PR for correctness. The AbstractSampler refactor can happen next. transition_save! can be fixed separately in AbstractMCMC as you mentioned. Then MixedVarInfo can be introduced and tested.

That said, most likely I will implement the above steps in this branch (I can close the PR though) and then move pieces to PRs after it's all working together. There might be more gotchas I find along the way that change the above plan so I don't want to commit too much (making and merging PRs) before it's all working locally.

[yebai]

@mohamed82008 Can we prioritise to finish this PR and/or related work in our NeurIPS submission? It's helpful to open a working PR and ask for feedback. It is ok to work on that slowly over the summer.

[yebai]

Closed in favour of #188