Remove context from model evaluation (use model.context instead)

[penelopeysm]

Summary

This PR modifies the model evaluation function, model.f for some model::DynamicPPL.Model, to not take a context as its third argument. Thus, its signature looks like f(model, varinfo, args...; kwargs...) where args and kwargs are forwarded from the function that defines the model.

During model evaluation, the tilde-pipeline would always dispatch on the __context__ argument, and would completely ignore the __model__.context. It has now been changed so that it dispatches on __model__.context.

As a result of this, we can remove the context argument from most model evaluation code as well as LogDensityFunction. This simplifies a lot of code.

There are a handful of minor follow-up points:

Combination of __context__ and __model__.context

Inside the model evaluation function, the model's context was previously being ignored. However, if one were to call evaluate!!(model, varinfo, context), the model's context does not get ignored: it is combined together with context in order to form a larger context stack, which then becomes the actual evaluation context.

DynamicPPL.jl/src/model.jl

Lines 942 to 944 in 3e54c2d

	context_new = setleafcontext(
	context, setleafcontext(model.context, leafcontext(context))
	)

After this PR, the intention is that if you want to keep that combining behaviour, you should do it manually. However, several 'outs' are provided:

1. sample!!(rng, model, varinfo, sampler) is provided. This wraps model.context in a SamplingContext(rng, sampler) before calling evaluate!!. This takes care of most invocations of evaluate!! where this combining behaviour was being used.

   (As a bonus, this also helps us remove the rng and sampler arguments from evaluate!!, which greatly simplifies that function: evaluate!! shenanigans #720)

2. evaluate!!(model, varinfo, context) still exists, but is dep-warned.

3. _evaluate!!(model, varinfo, context) still exists with the same behaviour without a depwarn. This is only retained because submodels depend on this and I would like to leave this change for a subsequent PR, as that needs to be reasoned about more deeply.

Models as callables

Previously, (model::Model)(args...) would forward to evaluate!!(args...). I think this behaviour is quite dangerous (#629) and even in that PR it was mentioned that we should be more explicit about what args we take, which I fully agree with.

Along with cleaning up evaluate!!, this PR also cleans up models as callables such that the only allowed signatures are model(), model(rng), model(varinfo), and model(rng, varinfo). Thus you are no longer allowed to specify a context (you should contextualize the model instead), or a sampler (the default sampler was SampleFromPrior() and given that model(...) was always being used to sample from the prior, it seems hugely unlikely that anybody was really passing a different sampler, and if they were, they can jolly well call first(DynamicPPL.sample!!(rng, model, varinfo, sampler))).

Remaining test failure

The remaining test failure (Julia-pre) is unrelated to this PR.

Closes #951

This is a step towards fixing #720 but it's not complete; that will have to wait for #960

[github-actions]

Benchmark Report for Commit fa90d1c

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 |                  8.5 |                 1.6 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                649.7 |                41.2 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                418.8 |                49.4 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |                986.3 |                35.9 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               6572.7 |                26.5 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1467.9 |                27.8 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |                989.9 |                 4.3 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5863.9 |                 3.8 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |                971.3 |                 9.0 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              66043.0 |                 3.5 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8883.8 |                 9.6 |
|               Dynamic |        10 |    mooncake |             typed |   true |                128.4 |                12.3 |
|              Submodel |         1 |    mooncake |             typed |   true |                 12.5 |                 6.2 |
|                   LDA |        12 | reversediff |             typed |   true |               1203.8 |                 2.6 |

[github-actions]

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

[codecov]

Codecov Report

Attention: Patch coverage is 89.31298% with 14 lines in your changes missing coverage. Please review.

Project coverage is 82.94%. Comparing base (bec523a) to head (fa90d1c).
Report is 1 commits behind head on breaking.

Files with missing lines	Patch %	Lines
src/simple_varinfo.jl	20.00%	8 Missing ⚠️
src/submodel_macro.jl	73.33%	4 Missing ⚠️
src/compiler.jl	66.66%	1 Missing ⚠️
src/test_utils/ad.jl	50.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking     #952      +/-   ##
============================================
+ Coverage     82.55%   82.94%   +0.38%     
============================================
  Files            38       38              
  Lines          4075     4068       -7     
============================================
+ Hits           3364     3374      +10     
+ Misses          711      694      -17     

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

That's an excellent direction. When duplicate contexts were first introduced, I was always bothered by them. We merged the relevant PRs to avoid blocking incremental improvements.

[yebai]

One minor comment on naming: sample!! will likely confuse people since it is ubiquitously used for sampling (posterior) distributions using MCMC. Perhaps consider an alternative, e.g., evaluate_and_sample!!?

[penelopeysm]

I'd be happy to rename. Will wait for @mhauru to review before doing it -- though I also realise that about 99% of usecases for sample!! exist because we're trying to construct a new varinfo using SampleFromPrior, and if we go through with #955 I think those would effectively turn into initialise!! rather than sample!! -- although that's probably for another time. (And either way, IMO we should keep it unexported)

[mhauru]

Comparing to this, this runtime seems to have ~doubled:

|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               6961.0 |                25.6 |

Not the end of the world, but any idea why? LDA has also gotten slower but I care less because it's so bad anyway.

[penelopeysm]

this runtime seems to have ~doubled

The time on this branch is the same as on the breaking branch, though: so maybe it came from an earlier merge into there? I'm just in the process of updating all those branches so can take a look in a while

[penelopeysm]

Also, while you're on leave, maybe this is a nice time for me to try to hack together something to monitor benchmark patterns over time haha

[mhauru]

My bad, I wasn't thinking when picking the thing to compare to. Nothing to do there then.

I'm still mid-review, but wondering about what to do about accumulators. I would like to say that I was hasty in merging the first accumulator stuff into breaking, because it's taking a while now to finish that stuff. It's releaseable in that tests pass, but I'm not sure it makes sense to release the accumulator stuff in parts. Thus we should move it to a different branch and clear the path for this to make it to a minor release. On the other hand though, this has been developed on top of accumulators, and I wouldn't want anyone to spend time disentangling the two with cherry-pick or rebase.

[mhauru]

I really like this PR. Just a few tiny typos and a couple of questions about code style.

[mhauru]

If there's an issue to track for this could link it. No worries if not.

[penelopeysm]

There are a few issues about callable structs but I couldn't find one that was explicitly about this. I'll try to make an MWE

[penelopeysm]

Okay, I fixed all the stuff. I renamed it to evaluate_and_sample!! too, but with a big warning in the changelog that this is liable to change (depending on what we do with SamplingContext). I think that should address the technical bits. The remaining question is whether we split this onto a separate branch. I'm inclined to not bother, partly because this isn't ready for release yet anyway (I need to fix it for submodels), and partly because of the hassle of cherry-picking (the simplification of leaf contexts helped a lot here, and I'm scared I might overlook correctness issues if I were to replay this on main).

I think the main argument for releasing it now is that I could use it upstream in some of the Turing sampler + LDF work, but I'm actually happy to hold off on that until we have a resolution to #955, and #955 in turn certainly needs to wait for accs to be done. So I think the happiest course of action is just for me to find other things to do until we are ready to release accs.

[mhauru]

Your argument for why it's okay to build this on top of accumulators is convincing and relieving.