Better handling of interaction between contexts

[torfjelde]

Things we can implement using contexts, i.e. use-cases to motivate better handling of "compositions" or "interplay" of contexts, i.e. PLEASE CAN WE REVISIT TuringLang/Turing.jl#254 OR SOME OTHER VERSION OF IT, I'M DYING OVER HERE.

Anyways:

• ConditionContext: allows specifying certain variables that are considered observations rather than assumptions, or even forcing observations to be considered missing/assumptions.
  • Following features are then trivial:
    • condition/decondition of a Model (if we introduce a ContextualModel which is just a pair of Context and Model), which in turn makes other contexts such as PriorContext and LikelihoodContext completely redundant.
    • Fixes Better warm up Turing.jl#273
    • Could be useful for Gibbs sampling?
• DistributionReplacementContext: allows replacement of distributions passed through the tilde-pipeline.
  • Following features are then trivial:
    • Automatic reparameterizations, e.g. somewhere in the pipeline we just replace the non-centered distribution with the centered one and apply the corresponding transformation on the resulting sample.
    • Performance-improvements. Often Distributions.jl's impl isn't the best, especially when combined with AD. Sure, we can just use @addlogprob!, etc., but this is annoying for users compared to just going optimize_distributions(model) or something and it just replaces distributions with nicer evaluators.
• SymbolicContext: allows using Symbolics.jl to trace through the model, thus opening up an entire new box of sweets! Working example is here ForwardDiff.gradient is slow Turing.jl#234, but it doesn't work well when used together with other contexts.
  • Following features are then trivial (for a lot of models):
    • DAG-extraction.
    • Symbolic generation of logdensity-computation.
    • Symbolic generation of gradient of logdensity-computation.
    • Simplifications.
    • And so much more.
• CUDAContext: makes it possible to execute the model on GPU. This is similar in nature to DistributionReplacementContext.
  • Following features are then trivial:
    • CUDA support (duh).

And all of these contexts can be hidden behind nicer functions a la condition.

So why do I have tears of sadness in my eyes you ask? Well, as things are currently, contexts do not work well together (I rave on and on about this in TuringLang/Turing.jl#249). E.g. if we introduced a ConditionContext, we want this to work well with MiniBatchContext, PrefixContext, etc., and also the more "lower-level" contexts such as DefaultContext and Turing.OptimizationContext. The same thing holds for SymbolicContext (which we also want to work in tandem with ConditionContext), CUDAContext, etc.

And you know the best part about all of this? Once we have some general handling of interactions between contexts, all of these contexts can be introduced independently of eachother without breaking any existing code! This makes it soooo much easier to go from "hey CUDA-support is a neat idea" to actually implementing it.

When we last spoke about this, the decision was to not do so at the time because the use-cases not clear enough + I mixed some ideas together, thus muddling the main-point: better handling of interactions between contexts makes introduction of new features much, much easier:)

So, I created this issue so we can rejuvenate the discussions + keep better track of what we're missing, why we want it, what are the issues, etc. rather than in a PR (as was done in TuringLang/Turing.jl#249).

[yebai]

Closed in favour of #895 and TuringLang/Turing.jl#744 -- we'd like to simplify context and restrict their interaction in principle.