InitContext, part 5 - Remove SamplingContext, SampleFrom{Prior,Uniform}, {tilde_,}assume (#985)

* Remove `SamplingContext` for good

* Remove `tilde_assume` as well

* Split up tilde_observe!! for Distribution / Submodel

* Tidy up tilde-pipeline methods and docstrings

* Fix tests

* fix ambiguity

* Add changelog

* Update HISTORY.md

Co-authored-by: Markus Hauru <[email protected]>

---------

Co-authored-by: Markus Hauru <[email protected]>

Author: penelopeysm

HISTORY.md

@@ -2,8 +2,66 @@
 
 ## 0.38.0
 
-The `varname_leaves` and `varname_and_value_leaves` functions have been moved to AbstractPPL.jl.
-Their behaviour is otherwise identical.
+**Breaking changes**
+
+### Introduction of `InitContext`
+
+DynamicPPL 0.38 introduces a new evaluation context, `InitContext`.
+It is used to generate fresh values for random variables in a model.
+
+Evaluation contexts are stored inside a `DynamicPPL.Model` object, and control what happens with tilde-statements when a model is run.
+The two major leaf (basic) contexts are `DefaultContext` and, now, `InitContext`.
+`DefaultContext` is the default context, and it simply uses the values that are already stored in the `VarInfo` object passed to the model evaluation function.
+On the other hand, `InitContext` ignores values in the VarInfo object and inserts new values obtained from a specified source.
+(It follows also that the VarInfo being used may be empty, which means that `InitContext` is now also the way to obtain a fresh VarInfo for a model.)
+
+DynamicPPL 0.38 provides three flavours of _initialisation strategies_, which are specified as the second argument to `InitContext`:
+
+  - `InitContext(rng, InitFromPrior())`: New values are sampled from the prior distribution (on the right-hand side of the tilde).
+  - `InitContext(rng, InitFromUniform(a, b))`: New values are sampled uniformly from the interval `[a, b]`, and then invlinked to the support of the distribution on the right-hand side of the tilde.
+  - `InitContext(rng, InitFromParams(p, fallback))`: New values are obtained by indexing into the `p` object, which can be a `NamedTuple` or `Dict{<:VarName}`. If a variable is not found in `p`, then the `fallback` strategy is used, which is simply another of these strategies. In particular, `InitFromParams` enables the case where different variables are to be initialised from different sources.
+
+(It is possible to define your own initialisation strategy; users who wish to do so are referred to the DynamicPPL API documentation and source code.)
+
+**The main impact on the upcoming Turing.jl release** is that, instead of providing initial values for sampling, the user will be expected to provide an initialisation strategy instead.
+This is a more flexible approach, and not only solves a number of pre-existing issues with initialisation of Turing models, but also improves the clarity of user code.
+In particular:
+
+  - When providing a set of fixed parameters (i.e. `InitFromParams(p)`), `p` must now either be a NamedTuple or a Dict. Previously Vectors were allowed, which is error-prone because the ordering of variables in a VarInfo is not obvious.
+  - The parameters in `p` must now always be provided in unlinked space (i.e., in the space of the distribution on the right-hand side of the tilde). Previously, whether a parameter was expected to be in linked or unlinked space depended on whether the VarInfo was linked or not, which was confusing.
+
+### Removal of `SamplingContext`
+
+For developers working on DynamicPPL, `InitContext` now completely replaces what used to be `SamplingContext`, `SampleFromPrior`, and `SampleFromUniform`.
+Evaluating a model with `SamplingContext(SampleFromPrior())` (e.g. with `DynamicPPL.evaluate_and_sample!!(model, VarInfo(), SampleFromPrior())` has a direct one-to-one replacement in `DynamicPPL.init!!(model, VarInfo(), InitFromPrior())`.
+Please see the docstring of `init!!` for more details.
+Likewise `SampleFromUniform()` can be replaced with `InitFromUniform()`.
+`InitFromParams()` provides new functionality which was previously implemented in the roundabout way of manipulating the VarInfo (e.g. using `unflatten`, or even more hackily by directly modifying values in the VarInfo), and then evaluating using `DefaultContext`.
+
+The main change that this is likely to create is for those who are implementing samplers or inference algorithms.
+The exact way in which this happens will be detailed in the Turing.jl changelog when a new release is made.
+Broadly speaking, though, `SamplingContext(MySampler())` will be removed so if your sampler needs custom behaviour with the tilde-pipeline you will likely have to define your own context.
+
+### Simplification of the tilde-pipeline
+
+There are now only two functions in the tilde-pipeline that need to be overloaded to change the behaviour of tilde-statements, namely, `tilde_assume!!` and `tilde_observe!!`.
+Other functions such as `tilde_assume` and `assume` (and their `observe` counterparts) have been removed.
+
+Note that this was effectively already the case in DynamicPPL 0.37 (where they were just wrappers around each other).
+The separation of these functions was primarily implemented to avoid performing extra work where unneeded (e.g. to not calculate the log-likelihood when `PriorContext` was being used). This functionality has since been replaced with accumulators (see the 0.37 changelog for more details).
+
+**Other changes**
+
+### Reimplementation of functions using `InitContext`
+
+A number of functions have been reimplemented and unified with the help of `InitContext`.
+In particular, this release brings substantial performance improvements for `returned` and `predict`.
+Their APIs are the same.
+
+### Upstreaming of VarName functionality
+
+The implementation of the `varname_leaves` and `varname_and_value_leaves` functions have been moved to AbstractPPL.jl.
+Their behaviour is otherwise identical, and they are still accessible from the DynamicPPL module (though still not exported).
 
 ## 0.37.3
 

docs/src/api.md

@@ -8,7 +8,7 @@ Part of the API of DynamicPPL is defined in the more lightweight interface packa
 
 A core component of DynamicPPL is the [`@model`](@ref) macro.
 It can be used to define probabilistic models in an intuitive way by specifying random variables and their distributions with `~` statements.
-These statements are rewritten by `@model` as calls of [internal functions](@ref model_internal) for sampling the variables and computing their log densities.
+These statements are rewritten by `@model` as calls of internal functions for sampling the variables and computing their log densities.
 
 ```@docs
 @model
@@ -344,6 +344,13 @@ Base.empty!
 SimpleVarInfo
 ```
 
+### Tilde-pipeline
+
+```@docs
+tilde_assume!!
+tilde_observe!!
+```
+
 ### Accumulators
 
 The subtypes of [`AbstractVarInfo`](@ref) store the cumulative log prior and log likelihood, and sometimes other variables that change during executing, in what are called accumulators.
@@ -447,12 +454,12 @@ AbstractPPL.evaluate!!
 
 This method mutates the `varinfo` used for execution.
 By default, it does not perform any actual sampling: it only evaluates the model using the values of the variables that are already in the `varinfo`.
+If you wish to sample new values, see the section on [VarInfo initialisation](#VarInfo-initialisation) just below this.
 
 The behaviour of a model execution can be changed with evaluation contexts, which are a field of the model.
 Contexts are subtypes of `AbstractPPL.AbstractContext`.
 
 ```@docs
-SamplingContext
 DefaultContext
 PrefixContext
 ConditionContext
@@ -486,15 +493,7 @@ DynamicPPL.init
 
 ### Samplers
 
-In DynamicPPL two samplers are defined that are used to initialize unobserved random variables:
-[`SampleFromPrior`](@ref) which samples from the prior distribution, and [`SampleFromUniform`](@ref) which samples from a uniform distribution.
-
-```@docs
-SampleFromPrior
-SampleFromUniform
-```
-
-Additionally, a generic sampler for inference is implemented.
+In DynamicPPL a generic sampler for inference is implemented.
 
 ```@docs
 Sampler
@@ -520,9 +519,3 @@ There is also the _experimental_ [`DynamicPPL.Experimental.determine_suitable_va
 DynamicPPL.Experimental.determine_suitable_varinfo
 DynamicPPL.Experimental.is_suitable_varinfo
 ```
-
-### [Model-Internal Functions](@id model_internal)
-
-```@docs
-tilde_assume
-```

ext/DynamicPPLEnzymeCoreExt.jl

@@ -8,8 +8,6 @@ else
     using ..EnzymeCore
 end
 
-@inline EnzymeCore.EnzymeRules.inactive_type(::Type{<:DynamicPPL.SamplingContext}) = true
-
 # Mark istrans as having 0 derivative. The `nothing` return value is not significant, Enzyme
 # only checks whether such a method exists, and never runs it.
 @inline EnzymeCore.EnzymeRules.inactive(::typeof(DynamicPPL.istrans), args...) = nothing

src/DynamicPPL.jl

@@ -96,18 +96,16 @@ export AbstractVarInfo,
     values_as_in_model,
     # Samplers
     Sampler,
-    SampleFromPrior,
-    SampleFromUniform,
     # LogDensityFunction
     LogDensityFunction,
     # Contexts
     contextualize,
-    SamplingContext,
     DefaultContext,
     PrefixContext,
     ConditionContext,
-    assume,
-    tilde_assume,
+    # Tilde pipeline
+    tilde_assume!!,
+    tilde_observe!!,
     # Initialisation
     InitContext,
     AbstractInitStrategy,