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)

I was too lazy to retroactively apply this fix to 0.39 (there's a lot of code that was changed and I honestly can't remember how it used to work). I think it's fine since we're close enough to merging 0.40.

[github-actions]

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

[codecov]

Codecov Report

❌ Patch coverage is 0% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 0.00%. Comparing base (d75e6f2) to head (61db751).

Files with missing lines	Patch %	Lines
src/mcmc/prior.jl	0.00%	8 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (d75e6f2) and HEAD (61db751). Click for more details.

HEAD has 46 uploads less than BASE

Flag	BASE (d75e6f2)	HEAD (61db751)
	48	2

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #2649       +/-   ##
==========================================
- Coverage   84.88%   0.00%   -84.89%     
==========================================
  Files          22      22               
  Lines        1475    1449       -26     
==========================================
- Hits         1252       0     -1252     
- Misses        223    1449     +1226     

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

Details

• 0 of 6 (0.0%) changed or added relevant lines in 1 file are covered.
• 1223 unchanged lines in 21 files lost coverage.
• Overall coverage decreased (-85.0%) to 0.0%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/mcmc/prior.jl	0	6	0.0%

Files with Coverage Reduction	New Missed Lines	%
src/mcmc/algorithm.jl	1	0.0%
src/mcmc/prior.jl	5	0.0%
src/Turing.jl	5	0.0%
src/mcmc/is.jl	15	0.0%
src/mcmc/abstractmcmc.jl	16	0.0%
src/stdlib/RandomMeasures.jl	22	0.0%
src/mcmc/repeat_sampler.jl	23	0.0%
ext/TuringDynamicHMCExt.jl	27	0.0%
src/mcmc/ess.jl	37	0.0%
src/mcmc/external_sampler.jl	42	0.0%

Totals
Change from base Build 16800936663:	-85.0%
Covered Lines:	0
Relevant Lines:	1449

---

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