`InitContext`, part 3 - Introduce `InitContext` #981

penelopeysm · 2025-07-09T23:32:58Z

Part 1: Adding hasvalue and getvalue to AbstractPPL
Part 2: Removing hasvalue and getvalue from DynamicPPL

This is Part 3/N of splitting up #967.

This PR solely adds the code for a new leaf context, InitContext. Its behaviour is to always override existing values inside a VarInfo.

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 this end, I've put all the InitContext-related code inside src/contexts/init.jl. This should, hopefully, also make reviewing easier.

InitContext comes in three forms:

PriorInit: Samples values from the prior;
UniformInit(min, max): Samples values from between min and max before invlinking them back into the support of the original distribution;
ParamsInit(params): Takes values from params, which can be a NamedTuple or a Dict. Also takes a fallback strategy, which is used in case the varname is not found in params.

PriorInit and UniformInit have almost one-to-one correspondence with SampleFromPrior and SampleFromUniform. I haven't removed them yet because a LOT of stuff depends on SampleFromPrior. We will eventually remove all of it but doing it in a single PR would make it way too big.

However, ParamsInit is new. Its purpose is to set fixed values in a VarInfo in a principled manner. Right now, there is a fair amount of code that does low-level VarInfo manipulation by reaching inside and modifying the underlying arrays, such as initialise_values!!, and setval_and_resample!. This is prone to bugs, and also leads to (for example) the VarInfo's logp being out of sync. (All of that code will be removed in subsequent PRs.)

ParamsInit accomplishes the same in a much cleaner way, although in some cases (specifically, when setting initial values for sampling) it may necessitate one extra model evaluation. Because initialisation is not meant to happen within a tight loop (by definition it happens once), I don't consider performance to be an important issue.

One might ask what's the difference between InitContext{ParamsInit}, ConditionContext, and FixedContext. Consider the model @model f() = x ~ Normal(). Evaluating this with each of the contexts gives the following behaviour:

	InitContext{ParamsInit}	ConditionContext	FixedContext
logpdf of `x ~ Normal()`	Prior	Likelihood	Ignored
`x` in VarInfo	Yes	No	No

In other words, InitContext behaves as if you had run the model and sampled a new value for x, but that value just so happened to be the one that you provided. In contrast, Condition and FixedContext fundamentally change the model that is being run in that x is no longer a parameter.

Closes #375.

github-actions · 2025-07-09T23:43:32Z

Benchmark Report for Commit `6705d7b`

Computer Information

Julia Version 1.11.6
Commit 9615af0f269 (2025-07-09 12:58 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 |                  8.8 |                 1.5 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                666.5 |                44.5 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                429.8 |                53.4 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1203.4 |                29.6 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               6758.9 |                29.4 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1064.0 |                40.1 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |               1015.7 |                 4.8 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5870.7 |                 4.3 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1019.0 |                 9.0 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              66219.4 |                 3.9 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8779.0 |                 9.9 |
|               Dynamic |        10 |    mooncake |             typed |   true |                142.6 |                11.4 |
|              Submodel |         1 |    mooncake |             typed |   true |                 13.0 |                 5.5 |
|                   LDA |        12 | reversediff |             typed |   true |               1065.9 |                 2.5 |

github-actions · 2025-07-10T00:13:43Z

DynamicPPL.jl documentation for PR #981 is available at:
https://TuringLang.github.io/DynamicPPL.jl/previews/PR981/

codecov · 2025-07-10T00:32:50Z

Codecov Report

❌ Patch coverage is 89.39394% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.53%. Comparing base (7b55aa3) to head (6705d7b).
⚠️ Report is 1 commits behind head on breaking.

Files with missing lines	Patch %	Lines
src/contexts/init.jl	92.30%	4 Missing ⚠️
src/model.jl	72.72%	3 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking     #981      +/-   ##
============================================
+ Coverage     82.26%   82.53%   +0.26%     
============================================
  Files            38       39       +1     
  Lines          3947     4008      +61     
============================================
+ Hits           3247     3308      +61     
  Misses          700      700

☔ 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.

penelopeysm · 2025-07-10T01:00:14Z

src/contexts.jl

-"""
-    prefix(model::Model, x::VarName)
-    prefix(model::Model, x::Val{sym})
-    prefix(model::Model, x::Any)
-


This code was shifted verbatim to src/model.jl to avoid circular dependencies between files.

Note that, apart from being simpler code, Distributions.Uniform also doesn't allow the lower and upper bounds to be exactly equal (but we might like to keep that option open in DynamicPPL, e.g. if the user wants to initialise all values to the same value in linked space).

This should have been changed in #940, but slipped through as the file wasn't listed as one of the changed files.

mhauru · 2025-08-12T15:17:10Z

src/contexts/init.jl

+!!! warning "Values must be unlinked"
+    The values returned by `init` are always in the untransformed space, i.e.,
+    they must be within the support of the original distribution. That means that,
+    for example, `init(rng, dist, u::UniformInit)` will in general return values that


What about calling it UniformLinkedInit to avoid a confusion here?

Hmmmm I don't have a super strong preference (the current docstring clarifies it, and the lack of a docstring was IMO the main problem with SampleFromUniform), but let's go with the more explicit and clearer name.

Although I also don't want to give the impression that init(..., ::UniformLinkedInit) returns values in linked space. :/ On this basis I might prefer the original UniformInit a bit better.

Hmm, fair point. I'm a bit torn between those two arguments, can leave as is if that's your preference.

I'll keep it as UniformInit for this PR but also check if others have an opinion, if necessary we can change in another PR before release.

Renamed to InitFromUniform, it's a bit longer but I agree with @yebai that it's clearer.

mhauru · 2025-08-12T15:17:45Z

src/contexts/init.jl

+    AbstractInitStrategy
+
+Abstract type representing the possible ways of initialising new values for
+the random variables in a model (e.g., when creating a new VarInfo).


Could this have a list of functions subtypes must implement methods for?

Good call, done.

mhauru · 2025-08-12T15:25:21Z

src/contexts/init.jl

+    return if hasvalue(p.params, vn, dist)
+        x = getvalue(p.params, vn, dist)
+        if x === missing
+            init(rng, vn, dist, p.default)


Could there be demand for p.default=nothing in which case this would error? I wonder if in some cases the current implementation could cause silent unexpected behaviour, e.g. if you misspell the varname and thus end up getting samples from the prior rather than the expected value.

Yup. I can't remember what prompted it now, but I did think about this a while back. I don't think that this would get used anywhere in DynamicPPL, but I'm not opposed to adding it in anyway.

Also, do you prefer p.default, or p.fallback?

Quite ambivalent between those too. I think the docs may have used one term and the code another.

I'm starting to like fallback more, because 'the default for default' feels weird, so if you're not opinionated I'll change it.

test/contexts.jl

Co-authored-by: Markus Hauru <markus@mhauru.org>

* Bump minor version * bump benchmarks compat * add a skeletal changelog * `InitContext`, part 3 - Introduce `InitContext` (#981) * Implement InitContext * Fix loading order of modules; move `prefix(::Model)` to model.jl * Add tests for InitContext behaviour * inline `rand(::Distributions.Uniform)` Note that, apart from being simpler code, Distributions.Uniform also doesn't allow the lower and upper bounds to be exactly equal (but we might like to keep that option open in DynamicPPL, e.g. if the user wants to initialise all values to the same value in linked space). * Document * Add a test to check that `init!!` doesn't change linking * Fix `push!` for VarNamedVector This should have been changed in #940, but slipped through as the file wasn't listed as one of the changed files. * Add some line breaks Co-authored-by: Markus Hauru <markus@mhauru.org> * Add the option of no fallback for ParamsInit * Improve docstrings * typo * `p.default` -> `p.fallback` * Rename `{Prior,Uniform,Params}Init` -> `InitFrom{Prior,Uniform,Params}` --------- Co-authored-by: Markus Hauru <markus@mhauru.org> * use `varname_leaves` from AbstractPPL instead (#1030) * use `varname_leaves` from AbstractPPL instead * add changelog entry * fix import * tidy occurrences of varname_leaves as well (#1031) * `InitContext`, part 4 - Use `init!!` to replace `evaluate_and_sample!!`, `predict`, `returned`, and `initialize_values` (#984) * Replace `evaluate_and_sample!!` -> `init!!` * Use `ParamsInit` for `predict`; remove `setval_and_resample!` and friends * Use `init!!` for initialisation * Paper over the `Sampling->Init` context stack (pending removal of SamplingContext) * Remove SamplingContext from JETExt to avoid triggering `Sampling->Init` pathway * Remove `predict` on vector of VarInfo * Fix some tests * Remove duplicated test * Simplify context testing * Rename FooInit -> InitFromFoo * Fix JETExt * Fix JETExt properly * Fix tests * Improve comments * Remove duplicated tests * Docstring improvements Co-authored-by: Markus Hauru <markus@mhauru.org> * Concretise `chain_sample_to_varname_dict` using chain value type * Clarify testset name * Re-add comment that shouldn't have vanished * Fix stale Requires dep * Fix default_varinfo/initialisation for odd models * Add comment to src/sampler.jl Co-authored-by: Markus Hauru <markus@mhauru.org> --------- Co-authored-by: Markus Hauru <markus@mhauru.org> * `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 <markus@mhauru.org> --------- Co-authored-by: Markus Hauru <markus@mhauru.org> * fix missing import * Shuffle context code around and remove dead code (#1050) * Delete the `"del"` flag (#1058) * Delete del * Fix a typo * Add HISTORY entry about del * Fixes for Turing 0.41 (#1057) * setleafcontext(model, ctx) and various other fixes * fix a bug * Add warning for `initial_parameters=...` * Remove `resume_from` and `default_chain_type` (#1061) * Remove resume_from * Format * Fix test * remove initial_params warning * Allow more flexible `initial_params` (#1064) * Enable NamedTuple/Dict initialisation * Add more tests * fix include_all kwarg for predict, improve perf (#1068) * Fix `include_all` for predict * Fix include_all for predict, some perf improvements * Replace `Metadata.flags` with `Metadata.trans` (#1060) * Replace Medata.flags with Metadata.trans * Fix a bug * Fix a typo * Fix two bugs * Rename trans to is_transformed * Rename islinked to is_transformed, remove duplication * Change pointwise_logdensities default key type to VarName (#1071) * Change pointwise_logdensities default key type to VarName * Fix a doctest * Fix DynamicPPL / MCMCChains methods (#1076) * Reimplement pointwise_logdensities (almost completely) * Move logjoint, logprior, ... as well * Fix imports, etc * Remove tests that are failing (yes I learnt this from Claude) * Changelog * logpdf * fix docstrings * allow dict output * changelog * fix some comments * fix tests * Fix more imports * Remove stray n Co-authored-by: Markus Hauru <markus@mhauru.org> * Expand `logprior`, `loglikelihood`, and `logjoint` docstrings --------- Co-authored-by: Markus Hauru <markus@mhauru.org> * Remove `Sampler` and its interface (#1037) * Remove `Sampler` and `initialstep` * Actually just remove the entire file * forgot one function * Move sampling test utils to Turing as well * Update changelog to correctly reflect changes * [skip ci] Make changelog headings more consistent --------- Co-authored-by: Markus Hauru <markus@mhauru.org>

github-actions bot assigned penelopeysm Jul 9, 2025

penelopeysm commented Jul 10, 2025

View reviewed changes

This was referenced Jul 10, 2025

The big InitContext PR. Do not merge! #967

Closed

InitContext, part 4 - Use init!! to replace evaluate_and_sample!!, predict, returned, and initialize_values #984

Merged

penelopeysm force-pushed the py/hasgetvalue branch from 90e95e1 to 21cf568 Compare July 10, 2025 11:59

penelopeysm changed the title ~~Introduce InitContext~~ InitContext, part 3 - Introduce InitContext Jul 10, 2025

penelopeysm force-pushed the py/init-prior-uniform branch 2 times, most recently from 025aa8b to b55c1e1 Compare July 10, 2025 14:24

Base automatically changed from py/hasgetvalue to breaking July 25, 2025 14:22

penelopeysm force-pushed the py/init-prior-uniform branch 2 times, most recently from 5f60c46 to 4408efb Compare July 26, 2025 17:45

penelopeysm marked this pull request as ready for review July 26, 2025 17:46

penelopeysm force-pushed the py/init-prior-uniform branch from 4408efb to 9c07727 Compare July 26, 2025 18:25

penelopeysm force-pushed the py/init-prior-uniform branch from 9c07727 to ef038c6 Compare August 5, 2025 11:52

Base automatically changed from breaking to main August 7, 2025 09:45

penelopeysm force-pushed the py/init-prior-uniform branch from ef038c6 to 5961ca9 Compare August 8, 2025 10:15

penelopeysm mentioned this pull request Aug 8, 2025

v0.38 #1018

Merged

6 tasks

penelopeysm force-pushed the py/init-prior-uniform branch from 5961ca9 to 0656487 Compare August 8, 2025 10:20

penelopeysm changed the base branch from main to breaking August 8, 2025 10:20

penelopeysm force-pushed the breaking branch from 37aecf8 to b1cdc2a Compare August 8, 2025 10:21

penelopeysm force-pushed the py/init-prior-uniform branch 2 times, most recently from fd78d42 to bfcdbb9 Compare August 10, 2025 13:33

penelopeysm added 2 commits August 10, 2025 14:44

Implement InitContext

3b87baa

Fix loading order of modules; move prefix(::Model) to model.jl

75f69b9

penelopeysm added 5 commits August 10, 2025 14:44

Add tests for InitContext behaviour

1988baa

Document

d24f966

Add a test to check that init!! doesn't change linking

42d74b6

Fix push! for VarNamedVector

ab3e8da

This should have been changed in #940, but slipped through as the file wasn't listed as one of the changed files.

penelopeysm force-pushed the py/init-prior-uniform branch from bfcdbb9 to ab3e8da Compare August 10, 2025 13:44

penelopeysm requested a review from mhauru August 12, 2025 15:05

mhauru reviewed Aug 12, 2025

View reviewed changes

penelopeysm and others added 3 commits August 12, 2025 18:07

Add some line breaks

be8a1b3

Co-authored-by: Markus Hauru <markus@mhauru.org>

Add the option of no fallback for ParamsInit

c40a193

Improve docstrings

bcfdd93

penelopeysm requested a review from mhauru August 12, 2025 17:38

typo

9626ae9

mhauru approved these changes Aug 13, 2025

View reviewed changes

penelopeysm added 4 commits August 13, 2025 12:04

p.default -> p.fallback

1e4ce9f

Merge branch 'main' into py/init-prior-uniform

8330d87

Merge branch 'breaking' into py/init-prior-uniform

deed931

Rename {Prior,Uniform,Params}Init -> InitFrom{Prior,Uniform,Params}

6705d7b

penelopeysm merged commit 991e825 into breaking Aug 13, 2025
20 of 21 checks passed

penelopeysm deleted the py/init-prior-uniform branch August 13, 2025 16:47

penelopeysm mentioned this pull request Oct 18, 2025

Add scale parameter to SampleFromUniform #375

Closed

InitContext, part 3 - Introduce InitContext #981

InitContext, part 3 - Introduce InitContext #981

Uh oh!

Conversation

penelopeysm commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmark Report for Commit 6705d7b

Computer Information

Benchmark Results

Uh oh!

github-actions bot commented Jul 10, 2025

Uh oh!

codecov bot commented Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

penelopeysm Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm Aug 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm Aug 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

`InitContext`, part 3 - Introduce `InitContext` #981

`InitContext`, part 3 - Introduce `InitContext` #981

penelopeysm commented Jul 9, 2025 •

edited

Loading

github-actions bot commented Jul 9, 2025 •

edited

Loading

Benchmark Report for Commit `6705d7b`

codecov bot commented Jul 10, 2025 •

edited

Loading

penelopeysm Jul 10, 2025 •

edited

Loading

penelopeysm Aug 12, 2025 •

edited

Loading

penelopeysm Aug 13, 2025 •

edited

Loading