Fix condition and fix in submodels

[penelopeysm]

Summary

This PR closes #857, i.e., gives users sensible ways of conditioning and fixing variables in submodels. To aid in understanding this, it also adds a long piece of documentation (PREVIEW HERE) explaining the design.

Performance

In the table below I compare the time taken for _evaluate!!(model, ...) on breaking (the base branch) and py/submodel-cond (HEAD of this PR). For good measure, the followup PR #896 is also benchmarked.

The model tested comprises m submodels, all of which contain n assumed variables. All times are in µs.

Profiling code (click to expand)

using DynamicPPL, Distributions, Chairmarks
using Plots

@model function inner(n)
    xs = Vector{Float64}(undef, n)
    for i in eachindex(xs)
        xs[i] ~ Normal(0, 1)
    end
end
@model function outer1(n)
    a ~ to_submodel(inner(n))
end
@model function outer10(n)
    a1 ~ to_submodel(inner(n))
    a2 ~ to_submodel(inner(n))
    a3 ~ to_submodel(inner(n))
    a4 ~ to_submodel(inner(n))
    a5 ~ to_submodel(inner(n))
    a6 ~ to_submodel(inner(n))
    a7 ~ to_submodel(inner(n))
    a8 ~ to_submodel(inner(n))
    a9 ~ to_submodel(inner(n))
    a10 ~ to_submodel(inner(n))
end
@model function outer20(n)
    a1 ~ to_submodel(inner(n))
    a2 ~ to_submodel(inner(n))
    a3 ~ to_submodel(inner(n))
    a4 ~ to_submodel(inner(n))
    a5 ~ to_submodel(inner(n))
    a6 ~ to_submodel(inner(n))
    a7 ~ to_submodel(inner(n))
    a8 ~ to_submodel(inner(n))
    a9 ~ to_submodel(inner(n))
    a10 ~ to_submodel(inner(n))
    a11 ~ to_submodel(inner(n))
    a12 ~ to_submodel(inner(n))
    a13 ~ to_submodel(inner(n))
    a14 ~ to_submodel(inner(n))
    a15 ~ to_submodel(inner(n))
    a16 ~ to_submodel(inner(n))
    a17 ~ to_submodel(inner(n))
    a18 ~ to_submodel(inner(n))
    a19 ~ to_submodel(inner(n))
    a20 ~ to_submodel(inner(n))
end

function profile(m, n)
    @info "Profiling with $m submodel(s) and $n inner model size"
    if m == 1
        model = outer1(n)
    elseif m == 10
        model = outer10(n)
    elseif m == 20
        model = outer20(n)
    else
        error("Invalid value for m")
    end
    v = VarInfo(model); c = DefaultContext();
    b = @be DynamicPPL._evaluate!!(model, v, c)
    @info "... got $(median(b).time)"
    return median(b).time
end
ms = [1, 10, 20]
ns = [1, 10, 25, 50, 100, 200]

# call profile(m, n) for m in ms for n in ns

         PR base      #892 - THIS PR    #896
m   n    breaking     py/submodel-cond  py/submodel-prefix
1   1    0.665634146  0.77410526        0.403409091
1   10   1.45625      1.60305263        1.16836
1   25   2.9          3.0375            2.549272727
1   50   4.9834       5.0418            4.645833333
1   100  9.416666667  9.43066667        9.3125
1   200  18.125       18.042            17.709
10  1    44.333       49                39.5
10  10   61.7085      65.458            54.625
10  25   85.875       92.771            81.458
10  50   136.458      140.812           129.042
10  100  235.708      232.6875          225.125
10  200  411.875      403.125           399.291
20  1    594.708      605.167           579.9585
20  10   650.9995     671.271           627.7915
20  25   734.417      759.937           708.375
20  50   866.292      876.521           839.708
20  100  1135.1665    1143.646          1110.6875
20  200  1702.291     1701.792          1644.166

TODO

• Check performance
  • Understand why performance is improved
• Apply the same fixes for fix
• Check whether the behaviour of decondition and unfix is still sensible (it should be, but you never know...)
• Document the approach taken in managing stacks of PrefixContext and (ConditionContext | FixedContext)
  • Correct the blurb about tilde_assume (it's no longer accurate after this PR)
  • Add a section on getting nested submodels right (i.e. PrefixContexts of PrefixContexts)
• Add changelog entry
• Add high-level tests for behaviour at the model level
  • paying special attention to nested submodels, because that's what broke [wip] submodel explorations #815;
  • also make sure to check manual prefixing and no prefixing;
  • also make sure to check the old-style conditioning via model arguments (the code I modified should not affect this because it doesn't touch the inargnames checks, only the contextual_isassumption checks; but probably good to test anyway)
  • also make sure to test with non-identity-lens varnames on LHS of submodel tilde
• Add unit tests for new functions

[codecov]

Codecov Report

Attention: Patch coverage is 97.61905% with 2 lines in your changes missing coverage. Please review.

Project coverage is 85.04%. Comparing base (be27636) to head (fcb44e5).
Report is 3 commits behind head on breaking.

Files with missing lines	Patch %	Lines
src/contexts.jl	98.30%	1 Missing ⚠️
src/utils.jl	80.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking     #892      +/-   ##
============================================
+ Coverage     84.76%   85.04%   +0.28%     
============================================
  Files            35       35              
  Lines          3879     3919      +40     
============================================
+ Hits           3288     3333      +45     
+ Misses          591      586       -5     

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

Mildly amazed that it passed all tests

I'm silly, I deleted the test suite, of course it passed

[github-actions]

Benchmark Report for Commit fcb44e5

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.1 |                 1.6 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                712.4 |                33.4 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                401.0 |                46.8 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1187.2 |                26.7 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               3305.2 |                22.5 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1410.5 |                29.2 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |                924.6 |                 5.1 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5241.9 |                 4.2 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |                976.0 |                 8.9 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              58686.3 |                 3.7 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8484.6 |                 9.8 |
|               Dynamic |        10 |    mooncake |             typed |   true |                125.0 |                11.9 |
|              Submodel |         1 |    mooncake |             typed |   true |                 13.1 |                 6.6 |
|                   LDA |        12 | reversediff |             typed |   true |                452.6 |                 5.3 |

[penelopeysm]

There's a nasty 4x regression on the submodel benchmark...

Submodels now actually run faster. I was prepared to take a slight hit on performance, but this has exceeded my expectations.

[penelopeysm]

@yebai Oh, I like the newsletter board idea!

[mhauru]

Looks really good, thanks for doing this. I only had tiny comments and questions.

In Turing.jl, GibbsContext might benefit from a similar treatment. Or replacing it with use of ConditionContext. There was some reason why we didn't do that originally, I forget now what it was.

[penelopeysm]

I have never taken a look at GibbsContext, but I reckon I can do so after this DPPL release.

[mhauru]

Thanks!