[WIP] InitContext

[penelopeysm]

This PR largely follows through on the proposal in #955. I have tried to keep the changes to as minimal a set as possible, but there is a lot of entwined code, sorry.

• Removes SamplingContext, SampleFromPrior, and SampleFromUniform

• Creates a new leaf context, InitContext. Its behaviour is to create new values, or overwrite existing variables in a VarInfo if they exist.

• Introduces initialisation strategies, PriorInit, UniformInit, and ParamsInit. See docstrings for more info.

• Changes evaluate_and_sample!! (introduced in Remove context from model evaluation (use model.context instead) #952) to init!!

• Moves all the VarInfo-initialisation code (used to be in sampler.jl) to ParamsInit.

Auxiliary changes that follow from the above:

• The function initialsampler, which was used to specify how initial values were generated for a given sampler, has been removed. The correct replacement is init_strategy(::AbstractSampler); however, that should be in Turing not DynamicPPL.

• Note that, after deleting the infamous assume method for SampleFromPrior, there is no longer any code that actually reads the del flag on VarInfo. There is a fair amount of code that writes to the del flag which has to be updated.

  • setval_and_resample! followed by model evaluation is replaced with ParamsInit with an incomplete set of parameters.
  • setval_and_resample! used to use nested_setindex_maybe. The role of this function was to handle cases where the VarName of interest was x, but instead of supplying x, the user supplied (say) x[1] and x[2]. This has now been replaced with hasvalue(dict, vn, dist) which uses dist to determine whether the sub-variables x[1] and x[2] provide enough information to reconstruct the full variable x. For example, this would return true if dist is MvNormal(zeros(2), I).

A long-held goal of mine is to split up contexts.jl and context_implementations.jl such that each context is in its own file. To help get us to this stage, I've put all the InitContext-related code inside src/contexts/init.jl. This should, hopefully, also make reviewing easier.

TODO for this PR

• ParamsInit: currently, initialising via NamedTuple or Dict is OK, but Vector will fail.
• Tests
• Docs (including proper user-facing docs, as I want InitStrategy to be exposed at Turing level, see "failed to find valid initial parameters" without the use of truncated Turing.jl#2476)

TODO for future PRs

• Split up contexts code

Look at this!:

Closes #859
Closes #955
Closes #797 (indirectly, since initialize_parameters is now init!!(..., ::ParamsInit))
Closes #774
Closes #375

Furthermore, once an upstream interface to InitContext is put into Turing, this will also solve:

TuringLang/Turing.jl#2476

[github-actions]

Benchmark Report for Commit 97df07f

Computer Information

Julia Version 1.11.5
Commit 760b2e5b739 (2025-04-14 06:53 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)

Benchmark Results

|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                  9.3 |                 1.6 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                693.4 |                38.7 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                433.3 |                49.0 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1005.3 |                35.3 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               6470.3 |                27.3 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1082.4 |                38.5 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |               1059.8 |                 4.3 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               6318.4 |                 3.6 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1012.1 |                 9.1 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              70058.2 |                 3.4 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8738.9 |                10.2 |
|               Dynamic |        10 |    mooncake |             typed |   true |                140.5 |                12.7 |
|              Submodel |         1 |    mooncake |             typed |   true |                 13.2 |                 6.5 |
|                   LDA |        12 | reversediff |             typed |   true |               1453.1 |                 5.0 |

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 156 lines in your changes missing coverage. Please review.

Project coverage is 4.29%. Comparing base (f20e86c) to head (97df07f).

Files with missing lines	Patch %	Lines
src/contexts/init.jl	0.00%	55 Missing ⚠️
src/test_utils/contexts.jl	0.00%	28 Missing ⚠️
src/utils.jl	0.00%	16 Missing ⚠️
src/model.jl	0.00%	15 Missing ⚠️
src/varinfo.jl	0.00%	15 Missing ⚠️
ext/DynamicPPLMCMCChainsExt.jl	0.00%	11 Missing ⚠️
src/simple_varinfo.jl	0.00%	10 Missing ⚠️
ext/DynamicPPLJETExt.jl	0.00%	2 Missing ⚠️
src/debug_utils.jl	0.00%	2 Missing ⚠️
src/extract_priors.jl	0.00%	1 Missing ⚠️
... and 1 more

❗ There is a different number of reports uploaded between BASE (f20e86c) and HEAD (97df07f). Click for more details.

HEAD has 6 uploads less than BASE

Flag	BASE (f20e86c)	HEAD (97df07f)
	12	6

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking    #967       +/-   ##
============================================
- Coverage     82.77%   4.29%   -78.48%     
============================================
  Files            38      39        +1     
  Lines          4011    3909      -102     
============================================
- Hits           3320     168     -3152     
- Misses          691    3741     +3050     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.