Use typed varinfo in Prior

[penelopeysm]

using Turing 

@model function f(x)
    θ ~ Normal()
    for i in eachindex(x)
        x[i] ~ Normal(θ)
    end
end

model = f(fill(missing, 100))

@time sample(model, Prior(), 1000; progress=false);

Base branch:

julia> @time sample(model, Prior(), 1000; progress=false);
  4.703977 seconds (19.72 M allocations: 718.042 MiB, 1.54% gc time)

This PR:

julia> @time sample(model, Prior(), 1000; progress=false);
  0.411107 seconds (5.65 M allocations: 242.001 MiB, 7.42% gc time)

[github-actions]

Turing.jl documentation for PR #2649 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2649/

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.64%. Comparing base (0a6a10b) to head (ec8f165).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2649      +/-   ##
==========================================
- Coverage   85.70%   85.64%   -0.06%     
==========================================
  Files          22       22              
  Lines        1434     1442       +8     
==========================================
+ Hits         1229     1235       +6     
- Misses        205      207       +2     

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

[coveralls]

Pull Request Test Coverage Report for Build 17377681902

Details

• 9 of 9 (100.0%) changed or added relevant lines in 1 file are covered.
• 2 unchanged lines in 1 file lost coverage.
• Overall coverage decreased (-0.05%) to 83.519%

---

Files with Coverage Reduction	New Missed Lines	%
src/mcmc/Inference.jl	2	84.16%

Totals
Change from base Build 17249540955:	-0.05%
Covered Lines:	1201
Relevant Lines:	1438

---

💛 - Coveralls

[mhauru]

Am I correct in thinking that this breaks prior sampling for

@model function f()
    x ~ Normal()
    if x > 0
        y ~ Normal()
    else
        z ~ Normal()
    end
end

but not for

@model function f()
    x ~ Normal()
    if x > 0
        y.a ~ Normal()
    else
        y.b ~ Normal()
    end
end

?

If yes, I'm happy.

[penelopeysm]

Actually, both are broken, for different reasons. The first one seems to never change, the second just doesn't run. I'll have to investigate this further, I'm a bit confused as to why either of them are broken.

[penelopeysm]

This works:

@model function f()
    x ~ Normal()
    if x > 0
        y ~ Normal()
    else
        z ~ Normal()
    end
end

It's because push!!(vi, ...) adds a new key to the NamedTuple-of-Metadata if it doesn't exist yet. https://github.com/TuringLang/DynamicPPL.jl/blob/f0ac1099c87326df95eada0aeb099029dbd09aee/src/varinfo.jl#L1723-L1735

MH and NUTS don't work on this model, only Prior does.

---

This doesn't work:

@model function f()
    x ~ Normal()
    y = (; a = 1.0, b = 1.0) # need this line so that it can assign to y.a or y.b
    if x > 0
        y.a ~ Normal()
    else
        y.b ~ Normal()
    end
end

The key y already exists in the NamedTuple, so it descends into push!(metadata, ...) which errors because metadata.idcs is a Dict which only has y.a and you can't push y.b into that dict unless you broaden the key type.

https://github.com/TuringLang/DynamicPPL.jl/blob/f0ac1099c87326df95eada0aeb099029dbd09aee/src/varinfo.jl#L1757-L1769

MH and NUTS also don't work on this model.

[mhauru]

I confirmed that both example models work on current release, when sampling with Prior.

I don't really mind losing

@model function f()
    x ~ Normal()
    if x > 0
        y ~ Normal()
    else
        z ~ Normal()
    end
end

I think that's syntax that we mostly don't support anyway.

The one mostly I care about is actually

@model function f()
    x ~ Normal()
    y = Vector{Float64}(undef, x > 0 ? 1 : 2)
    if x > 0
        y[1] ~ Normal()
    else
        y[1] ~ Normal()
        y[2] ~ Normal()
    end
end

That might be fine on this PR as-is?

I'm not sure if I care about

@model function f()
    x ~ Normal()
    y = (; a = 1.0, b = 1.0) # need this line so that it can assign to y.a or y.b
    if x > 0
        y.a ~ Normal()
    else
        y.b ~ Normal()
    end
end

I just didn't realise that it would behave differently from the Vector version when I proposed it as a test case.

[penelopeysm]

model with x, y[1], y[2] ... That might be fine on this PR as-is?

Yup, that still works on this PR. The x/y/z model also works on this PR, but x/y.a/y.b doesn't.

I think in terms of the code, this PR is mostly good to go. (I believe CI failures will be fixed by #2650.) But I'm a bit hesitant to merge into the upcoming minor version, because it does technically remove support for a few models.

In general, for arbitrary dynamic models like these, I think the solution to make sampling work is to use untyped VarInfo. That's what Prior used to do, which is also the only reason why it 'supported' all these things.

The long term solution would be to let people sample with whatever kind of varinfo they want. That was (is?) the purpose of the sample-with-LDF PRs I still have sitting around. But right now I'm kind of leaning towards not merging this until those are in, just in the interests of not gratuitously breaking things.

The flip side of course is that this PR does improve performance by a lot and it will probably be months until sample-with-LDF is merged.

[mhauru]

Without thinking this through carefully, I wonder if the performance gains are more important that losing support for some pretty niche models.

Actually, now that I said that, maybe those niche models include all models with dynamic use of submodels? Still kinda niche though.

Regardless, happy to leave it out of this upcoming release with DPPL 0.37.

[mhauru]

Is it time to revisit this with InitContext?

[penelopeysm]

Is it time to revisit this with InitContext?

I think InitContext makes things prettier but it doesn't solve the issue of same-symbol-but-different-type-VarName (y.a and y.b) model, and in particular it doesn't fix it when they crop up in different iterations, so

@model function f()
    if randn() > 0
        y.a ~ Normal()
    else
        y.b ~ Normal()
    end
end

There are a couple of ways of solving this properly:

1. Declare that we don't support such models (the assumption being that such models can be 'rescued' by renaming variables).
2. Fix push!! for typed VarInfo so that you can push y.b to a typed VarInfo that already has y.a.

I don't have a strong opinion on (1); but it does seem to me that independently of (1) we should still fix (2). To this end, I'll open an issue on DynamicPPL and also cheekily stick this on the list of topics for the JuliaCon hackathon.