Make PrefixContext contain a varname rather than symbol

[penelopeysm]

Warning

This PR should only be merged after #892

This PR changes PrefixContext to take a VarName as its 'prefix' rather than just a symbol. This allows PrefixContext to carry more information and therefore lets us correctly prefix variables in submodels: for example,

@model function inner()
    return x ~ Normal()
end
@model function outer()
    a = Vector{Float64}(undef, 1)
    a[1] ~ to_submodel(inner())
    return a
end
only(keys(VarInfo(outer()))) == @varname(a[1].x) # true

# Before this PR (i.e. on the `breaking` branch):
# only(keys(VarInfo(outer()))) == @varname(var"a[1]".x)

# v0.35 (i.e. on `main`):
# only(keys(VarInfo(outer()))) == @varname(var"a[1].x")

Note that this change does not affect type stability, because all information about a VarName (its symbol and its optic) is contained in its type.

The improvement to the prefixing behaviour follows from TuringLang/AbstractPPL.jl#119 and #830

Performance

Note

The below is largely copied from #892

In the table below I compare the time taken for _evaluate!!(model, ...) for this PR, the current breaking branch, and #892. 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

                      #892 - PR base    #896 - THIS PR
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

[github-actions]

Benchmark Report for Commit e866762

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.9 |                 1.5 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                753.7 |                33.9 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                414.4 |                47.9 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1227.5 |                26.9 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               3241.4 |                23.8 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1445.5 |                29.5 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |                934.5 |                 5.3 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5440.6 |                 4.1 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1008.8 |                 8.9 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              60949.3 |                 3.8 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               9441.9 |                 9.3 |
|               Dynamic |        10 |    mooncake |             typed |   true |                129.2 |                13.7 |
|              Submodel |         1 |    mooncake |             typed |   true |                 14.0 |                 6.4 |
|                   LDA |        12 | reversediff |             typed |   true |                473.5 |                 5.3 |

[codecov]

Codecov Report

Attention: Patch coverage is 95.65217% with 1 line in your changes missing coverage. Please review.

Project coverage is 85.04%. Comparing base (b545a93) to head (e866762).
Report is 3 commits behind head on py/submodel-cond.

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

Additional details and impacted files

@@                 Coverage Diff                  @@
##           py/submodel-cond     #896      +/-   ##
====================================================
- Coverage             85.05%   85.04%   -0.02%     
====================================================
  Files                    35       35              
  Lines                  3915     3919       +4     
====================================================
+ Hits                   3330     3333       +3     
- Misses                  585      586       +1     

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

[yebai]

Is #794 related?

[penelopeysm]

Is #794 related?

No, unfortunately not, I checked.

[mhauru]

I can't find a single thing to complain about, this is great!