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

penelopeysm · 2025-07-10T11:42:11Z

Part 1: Adding hasvalue and getvalue to AbstractPPL
Part 2: Removing hasvalue and getvalue from DynamicPPL
Part 3: Introducing InitContext and init!!

This is part 4/N of #967.

In Part 3 we introduced InitContext. This PR makes use of the functionality in there to replace a bunch of code that no longer needs to exist:

setval_and_resample! followed by model evaluation: This process was used for predict and returned, to manually store certain values in the VarInfo, which would be used in the subsequent model evaluation. We can now do this in a single step using InitFromParams.
initialize_values!!: very similar to the above. It would manually set values inside the varinfo, and then it would trigger an extra model evaluation to update the logp field. Again, this is directly replaced with InitFromParams.
evaluate_and_sample!!: direct one-to-one replacement with init!!.

There is one fairly major API change associated with point (2): the initial_params kwarg to Turing's sample must now be an AbstractInitStrategy.

It's still optional (it will default to init_strategy(spl), which is usually InitFromPrior, except for the HMC family which uses InitFromUniform). However, there are two implications:

initial_params cannot be a vector of parameters anymore. It must be InitFromParams(::NamedTuple) OR InitFromParams(::AbstractDict{VarName}).
Because InitFromParams expects values in unlinked space, initial_params must always be specified in unlinked space. Previously, initial_params would have to be specified in a way that matched the linking status of the underlying varinfo.

I consider both of these to be a major win for clarity. (One might argue that vectors are more convenient. But IMO anything that lets you extract a vector will also let you extract a NT or Dict, maybe with a bit more typing at worst).

Closes

Closes #774
Closes #797
Closes #983
Closes TuringLang/Turing.jl#2476
Closes TuringLang/Turing.jl#1775

github-actions · 2025-07-10T11:52:27Z

Benchmark Report for Commit `956ed54`

Computer Information

Julia Version 1.11.7
Commit f2b3dbda30a (2025-09-08 12:10 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.3 |                 1.6 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                654.6 |                45.1 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                426.3 |                55.7 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1186.8 |                30.0 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               6616.9 |                30.1 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1303.9 |                32.2 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |               1006.7 |                 4.8 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5775.9 |                 4.4 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1033.5 |                 8.7 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              65221.0 |                 4.0 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8885.5 |                 9.7 |
|               Dynamic |        10 |    mooncake |             typed |   true |                144.7 |                10.7 |
|              Submodel |         1 |    mooncake |             typed |   true |                 13.6 |                 4.9 |
|                   LDA |        12 | reversediff |             typed |   true |               1172.4 |                 3.1 |

github-actions · 2025-07-20T11:29:17Z

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

codecov · 2025-07-20T11:54:15Z

Codecov Report

❌ Patch coverage is 86.95652% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.20%. Comparing base (1e1cd94) to head (956ed54).

Files with missing lines	Patch %	Lines
src/simple_varinfo.jl	40.00%	6 Missing ⚠️
src/test_utils/contexts.jl	83.33%	5 Missing ⚠️
src/test_utils/model_interface.jl	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking     #984      +/-   ##
============================================
- Coverage     82.87%   81.20%   -1.68%     
============================================
  Files            39       39              
  Lines          3971     3910      -61     
============================================
- Hits           3291     3175     -116     
- Misses          680      735      +55

☔ 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-08-13T17:17:18Z

ext/DynamicPPLJETExt.jl

 function DynamicPPL.Experimental._determine_varinfo_jet(
    model::DynamicPPL.Model; only_ddpl::Bool=true
 )


I'm a bit confused by the comments in this function because as far as I can tell it only ever tested sampling, not both sampling and evaluation. (That was also true going further back e.g. in v0.36)

This PR thus also changes the implementation of this function to test both evaluation and sampling (i.e. initialisation) and if either fails, it will return the untyped varinfo.

Sorry I had to make this change in this PR. There were a few unholy tests where one would end up evaluating a model with a SamplingContext{<:InitContext}, which would error unless I introduced special code to handle it, and I didn't really want to do that. JETExt was one of those unholy scenarios.

penelopeysm · 2025-08-13T23:10:05Z

src/sampler.jl

-    DynamicPPL.evaluate_and_sample!!(rng, model, vi, sampler)
-    return vi, nothing
+    strategy = sampler isa SampleFromPrior ? InitFromPrior() : InitFromUniform()
+    _, new_vi = DynamicPPL.init!!(rng, model, vi, strategy)
+    return new_vi, nothing


This is a bit weird, but it's really just to tide us over until we delete SampleFromUniform/SampleFromPrior properly.

penelopeysm · 2025-08-13T23:11:42Z

src/sampler.jl

+Define the initialisation strategy used for generating initial values when
+sampling with `sampler`. Defaults to `InitFromPrior()`, but can be overridden.
+"""
+init_strategy(::Sampler) = InitFromPrior()


Eventually the aim would be to use ::AbstractSampler. But that will have to wait for cleanup in Turing. DynamicPPL itself doesn't use Sampler at all and if you only look at DPPL it looks like a meaningless empty wrapper, but Turing relies on these methods a fair bit

penelopeysm · 2025-08-13T23:19:25Z

test/model.jl

-    @testset "rng" begin
-        model = GDEMO_DEFAULT
-
-        for sampler in (SampleFromPrior(), SampleFromUniform())
-            for i in 1:10
-                Random.seed!(100 + i)
-                vi = VarInfo()
-                DynamicPPL.evaluate_and_sample!!(Random.default_rng(), model, vi, sampler)
-                vals = vi[:]
-
-                Random.seed!(100 + i)
-                vi = VarInfo()
-                DynamicPPL.evaluate_and_sample!!(Random.default_rng(), model, vi, sampler)
-                @test vi[:] == vals
-            end
-        end
-    end


I think this is adequately tested in the InitContext tests (test_rng_respected)

penelopeysm · 2025-08-13T23:21:06Z

src/test_utils/contexts.jl

-    varinfo_untyped = DynamicPPL.VarInfo()
-    model_with_spl = contextualize(model, SamplingContext(context))


I also had to rework these tests because of the SamplingContext{<:InitContext} case.

test/varinfo.jl

penelopeysm · 2025-08-13T23:25:35Z

test/varinfo.jl

        @test merge(vi_double, vi_single)[vn] == 1.0
    end

-    @testset "sampling from linked varinfo" begin


these tests are also covered in InitContext now (test_link_status_respected)!

mhauru

Pretty excited for this! Had some various nits and questions though.

src/model.jl

src/sampler.jl

src/test_utils/contexts.jl

test/compiler.jl

test/model.jl

test/test_util.jl

ext/DynamicPPLMCMCChainsExt.jl

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

src/sampler.jl

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

github-actions bot assigned penelopeysm Jul 10, 2025

penelopeysm mentioned this pull request Jul 10, 2025

The big InitContext PR. Do not merge! #967

Closed

4 tasks

penelopeysm changed the title ~~Use init!! to replace evaluate_and_sample!!, predict, returned, and initialize_values~~ InitContext, part 4 - Use init!! to replace evaluate_and_sample!!, predict, returned, and initialize_values 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

penelopeysm force-pushed the py/actually-use-init branch 5 times, most recently from b72c3bf to 92d3542 Compare July 10, 2025 15:57

penelopeysm mentioned this pull request Jul 10, 2025

Release v0.37 #901

Merged

22 tasks

penelopeysm force-pushed the py/actually-use-init branch 4 times, most recently from 7438b23 to d55d378 Compare July 10, 2025 16:56

penelopeysm mentioned this pull request Jul 10, 2025

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

Merged

penelopeysm force-pushed the py/actually-use-init branch 3 times, most recently from 12d93e5 to 7a8e7e3 Compare July 10, 2025 17:47

penelopeysm force-pushed the py/actually-use-init branch 2 times, most recently from 1d8bceb to 2edcd10 Compare July 20, 2025 00:59

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

penelopeysm force-pushed the py/actually-use-init branch from 5025592 to 3a16f9c Compare July 26, 2025 18:34

penelopeysm mentioned this pull request Jul 26, 2025

predict on vector of varinfo is incorrect if the varinfo passed in is linked #983

Open

penelopeysm force-pushed the py/actually-use-init branch from 3a16f9c to 4c96020 Compare July 26, 2025 18:47

penelopeysm marked this pull request as ready for review August 13, 2025 16:57

Fix JETExt

70bb2c4

penelopeysm force-pushed the py/actually-use-init branch from 89bc0ea to 726d486 Compare August 13, 2025 17:15

Fix JETExt properly

bc04355

penelopeysm force-pushed the py/actually-use-init branch from 726d486 to bc04355 Compare August 13, 2025 17:16

penelopeysm commented Aug 13, 2025

View reviewed changes

Fix tests

2cfc297

penelopeysm commented Aug 13, 2025

View reviewed changes

Improve comments

891b4b3

penelopeysm commented Aug 13, 2025

View reviewed changes

test/varinfo.jl Show resolved Hide resolved

Remove duplicated tests

3bb7ade

penelopeysm commented Aug 13, 2025

View reviewed changes

penelopeysm requested a review from mhauru August 13, 2025 23:54

Merge branch 'breaking' into py/actually-use-init

1bdb76e

mhauru requested changes Sep 16, 2025

View reviewed changes

penelopeysm and others added 5 commits September 16, 2025 10:14

Docstring improvements

39b958d

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

Concretise chain_sample_to_varname_dict using chain value type

907d24f

Clarify testset name

07946a7

Re-add comment that shouldn't have vanished

afdb173

Fix stale Requires dep

c641923

penelopeysm mentioned this pull request Sep 16, 2025

test more varinfos with parent contexts #1041

Open

Fix default_varinfo/initialisation for odd models

956ed54

penelopeysm mentioned this pull request Sep 17, 2025

improve implementation of predict(...; include_all) #1042

Open

mhauru approved these changes Sep 18, 2025

View reviewed changes

src/sampler.jl Show resolved Hide resolved

Add comment to src/sampler.jl

8d13c30

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

penelopeysm merged commit 729bfba into breaking Sep 18, 2025
3 of 19 checks passed

penelopeysm deleted the py/actually-use-init branch September 18, 2025 09:29

		varinfo_untyped = DynamicPPL.VarInfo()
		model_with_spl = contextualize(model, SamplingContext(context))

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

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

Uh oh!

Conversation

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

Uh oh!

Closes

Uh oh!

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

Uh oh!

Benchmark Report for Commit 956ed54

Computer Information

Benchmark Results

Uh oh!

github-actions bot commented Jul 20, 2025

Uh oh!

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

Uh oh!

Codecov Report

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!

penelopeysm Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

penelopeysm Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

penelopeysm Aug 13, 2025

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!

Uh oh!

penelopeysm Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

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

penelopeysm commented Jul 10, 2025 •

edited

Loading

github-actions bot commented Jul 10, 2025 •

edited

Loading

Benchmark Report for Commit `956ed54`

codecov bot commented Jul 20, 2025 •

edited

Loading

penelopeysm Aug 13, 2025 •

edited

Loading

penelopeysm Aug 13, 2025 •

edited

Loading