Fix potential undefined variable in model compiler output

[penelopeysm]

Fixes TuringLang/Libtask.jl#200.

What does this change?

To see what's different, consider the output of

@macroexpand @model function f(x)
    b ~ Normal()
    y ~ MvNormal([b], I)
end

Before this PR, this gives something like this for the y ~ MvNormal(...) line:

if isfixed(...)
elseif isassumption(...)
    ...
else
    if !((DynamicPPL.inargnames)(@varname(y), __model__))
        y = (DynamicPPL.getconditioned_nested)(__model__.context, (prefix)(__model__.context, @varname(y)))
    end
    DynamicPPL.tilde_observe!!(..., DynamicPPL.check_tilde_rhs(...), ..., y, ...)
    ...
end

This PR changes it to something like this:

if isfixed(...)
elseif isassumption(...)
    ...
else
    y = if !((DynamicPPL.inargnames)(@varname(y), __model__))
        (DynamicPPL.getconditioned_nested)(__model__.context, (prefix)(__model__.context, @varname(y)))
    else
        y
    end
    DynamicPPL.tilde_observe!!(..., DynamicPPL.check_tilde_rhs(...), ..., y, ...)
    ...
end

From DynamicPPL's point of view this is entirely equivalent, and indeed both model functions work perfectly fine, if you execute them one time from front to back.

Why?

The problem (as seen in TuringLang/Libtask.jl#200) comes with two things:

Firstly, when the Julia code in the first example is lowered to SSA IR, the compiler inserts a check right before tilde_observe!! which throws an error if !inargnames(...) was true and if y is undefined. I don't know exactly why this is needed, but it makes sense that if the y = DynamicPPL.getconditioned... line was executed, then y must be defined.

The IR looks something like this. The check is in the second-last line, and the value of %58 is (ultimately) derived from %52, which is !inargnames(...).

  25 ─ %52 = (DynamicPPL.inargnames)(%22, _2)::Any
  │    %53 = !%52::Any
  └───       goto #27 if not %53
  26 ─ %55 = Base.getfield(_2, :context)::DynamicPPL.ConditionContext{..}
  └─── %56 = (DynamicPPL.getconditioned_nested)(%55, %22)::Any
  27 ┄ %57 = φ (#26 => %56, #25 => #undef)::Any
  │    %58 = φ (#26 => true, #25 => false)::Bool
  │    %59 = Base.getfield(_2, :context)::DynamicPPL.ConditionContext{...}
  │    %60 = (DynamicPPL.check_tilde_rhs)(%14)::Any
  │          $(Expr(:throw_undef_if_not, :y, :(%58)))::Any
  │    %62 = (DynamicPPL.tilde_observe!!)(%59, %60, %57, %22, %12)::Tuple{Any, Any}

Secondly, when Libtask makes a resumable task out of this, it doesn't know whether check_tilde_rhs "produces". (See also footnote.) That means that it has to insert a resume point between the check_tilde_rhs and the call to :throw_undef_if_not.

The problem with this is, if check_tilde_rhs indeed produces, then when we resume the function, we'll jump straight to the throw_undef_if_not call. But if we jump in here, %58 is not defined, because we haven't run the inargnames check!

With this PR, it ensures that y is always defined regardless of whether it's inargnames or not. That means that the compiler doesn't insert the throw_undef_if_not line into the IR, and it means that Libtask doesn't fall over.

Notes

1. It seems odd to me that Libtask checks whether check_tilde_rhs returns a ProducedValue. I thought that this was opt-in, i.e. by default it should assume that check_tilde_rhs doesn't return a ProducedValue unless we declared it with might_produce. But this requires more knowledge of Libtask internals than I have right now.

2. Identifying the root cause also allows us to construct a Turing-free MWE for the same issue. One could argue that this is really an upstream bug, but fixing that will probably be substantially more difficult (and perhaps impossible).

3. In either case, I think it's just a good idea to fix it in DynamicPPL anyway because having variables that are only sometimes defined is not good (JET has also started to complain about these more often). I have a suspicion that the compiler's throw_undef_if_not is precisely what stops JET from complaining about it in this case (since we do run JET tests on model outputs).

4. One might wonder why Libtask only errored with specific models. I think that it's because for simpler models, the inargnames check (or possibly something at an even higher level, like isassumption) can be completely compiled away and so the emitted IR doesn't have this structure.

[github-actions]

Benchmark Report for Commit d528e3b

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 │   Dim │  AD Backend │           VarInfo │ Linked │ t(eval)/t(ref) │ t(grad)/t(eval) │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼────────────────┼─────────────────┤
│ Simple assume observe │     1 │ forwarddiff │             typed │  false │            6.6 │             1.7 │
│           Smorgasbord │   201 │ forwarddiff │             typed │  false │          729.7 │            44.0 │
│           Smorgasbord │   201 │ forwarddiff │ simple_namedtuple │   true │          414.2 │            55.8 │
│           Smorgasbord │   201 │ forwarddiff │           untyped │   true │          789.7 │            36.2 │
│           Smorgasbord │   201 │ forwarddiff │       simple_dict │   true │         7127.1 │            27.8 │
│           Smorgasbord │   201 │ forwarddiff │      typed_vector │   true │          762.4 │            43.5 │
│           Smorgasbord │   201 │ forwarddiff │    untyped_vector │   true │          775.9 │            36.9 │
│           Smorgasbord │   201 │ reversediff │             typed │   true │          916.9 │            44.4 │
│           Smorgasbord │   201 │    mooncake │             typed │   true │          740.0 │             5.7 │
│           Smorgasbord │   201 │      enzyme │             typed │   true │          915.3 │             3.7 │
│    Loop univariate 1k │  1000 │    mooncake │             typed │   true │         3844.1 │             5.8 │
│       Multivariate 1k │  1000 │    mooncake │             typed │   true │          974.4 │             9.1 │
│   Loop univariate 10k │ 10000 │    mooncake │             typed │   true │        42451.7 │             5.3 │
│      Multivariate 10k │ 10000 │    mooncake │             typed │   true │         8901.2 │             9.7 │
│               Dynamic │    10 │    mooncake │             typed │   true │          120.2 │            11.0 │
│              Submodel │     1 │    mooncake │             typed │   true │            8.6 │             6.4 │
│                   LDA │    12 │ reversediff │             typed │   true │          996.8 │             2.1 │
└───────────────────────┴───────┴─────────────┴───────────────────┴────────┴────────────────┴─────────────────┘

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.23%. Comparing base (d9af270) to head (d528e3b).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1110   +/-   ##
=======================================
  Coverage   81.23%   81.23%           
=======================================
  Files          40       40           
  Lines        3805     3805           
=======================================
  Hits         3091     3091           
  Misses        714      714           

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

I thought I'd add tests, but it seems that it's also impossible to reproduce this without ParticleMCMCContext or something similar that's complex enough, because the following

@model function f(x)
   β ~ Normal()
   μ = β .* x
   y ~ MvNormal(μ, I)
end
m = f([0.0]) | (; y=[0.0])
v = VarInfo(m)
args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(m, v)
Libtask.TapedTask(nothing, m.f, args...)

actually works just fine, even on current DPPL main, because all of the offending IR just gets compiled away. So I think the testing has to be done in Turing.

[github-actions]

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

[mhauru]

Looks good, thanks. I don't mind the lack of a test too much, it's hard to test for and the comment guards against someone reverting this inadvertently.

It seems odd to me that Libtask checks whether check_tilde_rhs returns a ProducedValue. I thought that this was opt-in, i.e. by default it should assume that check_tilde_rhs doesn't return a ProducedValue unless we declared it with might_produce. But this requires more knowledge of Libtask internals than I have right now.

I wonder if the answer to this has something to do with the distinction between :dynamic, :call, and :invoke, and which ones trigger the might_produce stuff. I couldn't explain it to you though.