fix include_all kwarg for predict, improve perf

[penelopeysm]

this PR:

• makes predict(model, chn; include_all) work: previously it would error unless the rng argument was also specified
• cuts out one model evaluation per MCMC iteration when predicting, which should make things nicely faster

benchmarks using the example from DPPL test suite:

using DynamicPPL, Distributions, MCMCChains, Chairmarks

@model function linear_reg(x, y, σ=0.1)
    β ~ Normal(0, 1)
    for i in eachindex(y)
        y[i] ~ Normal(β * x[i], σ)
    end
    # Insert a := block to test that it is not included in predictions
    return σ2 := σ^2
end

# Construct a chain with 'sampled values' of β
ground_truth_β = 2
β_chain = MCMCChains.Chains(
    rand(Normal(ground_truth_β, 0.002), 1000),
    [:β];
    info=(; varname_to_symbol=Dict(@varname(β) => :β)),
)

# Generate predictions from that chain
xs_test = [10 + 0.1, 10 + 2 * 0.1]
m_lin_reg_test = linear_reg(xs_test, fill(missing, length(xs_test)))
@be DynamicPPL.predict(m_lin_reg_test, β_chain)

# breaking
julia> @be DynamicPPL.predict(m_lin_reg_test, β_chain)
Benchmark: 5 samples with 1 evaluation
 min    19.908 ms (284965 allocs: 12.245 MiB)
 median 20.093 ms (284965 allocs: 12.245 MiB)
 mean   20.468 ms (285165 allocs: 12.453 MiB)
 max    22.045 ms (285965 allocs: 13.283 MiB)

# this PR
julia> @be DynamicPPL.predict(m_lin_reg_test, β_chain)
Benchmark: 8 samples with 1 evaluation
 min    12.631 ms (194998 allocs: 6.853 MiB)
 median 12.656 ms (194998 allocs: 6.853 MiB)
 mean   12.778 ms (194998 allocs: 6.853 MiB)
 max    13.553 ms (194998 allocs: 6.853 MiB)

[github-actions]

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.47%. Comparing base (ec65b4f) to head (b0f1e96).
⚠️ Report is 1 commits behind head on breaking.

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking    #1068      +/-   ##
============================================
+ Coverage     82.38%   82.47%   +0.09%     
============================================
  Files            42       42              
  Lines          3820     3824       +4     
============================================
+ Hits           3147     3154       +7     
+ Misses          673      670       -3     

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

[sunxd3]

LG

[penelopeysm]

If this is of any interest: Generally for small models MCMCChains does better than FlexiChains (I think this is mainly because constructing a FlexiChain from an array-of-dicts is quite slow because the internal storage is dict-of-array), but in the case where there are parameters that are very long vectors, MCMCChains scales very poorly because it has to split them up.

Note that this doesn't apply to the model in the top comment, since in that model, each element of y is a separate varname (so regardless of which chain backend is used, it gets split up).

### SETUP - using this PR
using DynamicPPL, Distributions, DistributionsAD, FlexiChains, MCMCChains, Chairmarks

@model function f(N, σ=0.1)
    m ~ Normal(0, 1)
    y ~ filldist(Normal(m, 1), N)
end

# Construct a chain with 'sampled values' of m
ground_truth_m = 2
randvals = rand(Normal(ground_truth_m, 0.002), 1000)
mchain = MCMCChains.Chains(
    randvals,
    [:m];
    info=(; varname_to_symbol=Dict(@varname(m) => :m)),
)
fchain = FlexiChains.FlexiChain{VarName}(
    1000, 1,
    Dict(FlexiChains.Parameter(@varname(m)) => randvals)
)

m1 = f(1)
@be DynamicPPL.predict(m1, fchain; include_all=false)
@be DynamicPPL.predict(m1, mchain; include_all=false)

m1000 = f(1000)
@be DynamicPPL.predict(m1000, fchain; include_all=false)
@be DynamicPPL.predict(m1000, mchain; include_all=false)

### RESULTS

julia> @be DynamicPPL.predict(m1, fchain; include_all=false)
Benchmark: 7 samples with 1 evaluation
 min    14.465 ms (193039 allocs: 8.451 MiB)
 median 15.053 ms (193039 allocs: 8.451 MiB)
 mean   16.036 ms (193039 allocs: 8.451 MiB, 5.14% gc time)
 max    22.861 ms (193039 allocs: 8.451 MiB, 35.95% gc time)

julia> @be DynamicPPL.predict(m1, mchain; include_all=false)
Benchmark: 9 samples with 1 evaluation
 min    11.015 ms (159456 allocs: 6.226 MiB)
 median 11.365 ms (159456 allocs: 6.226 MiB)
 mean   12.149 ms (159456 allocs: 6.226 MiB, 4.55% gc time)
 max    18.945 ms (159456 allocs: 6.226 MiB, 40.98% gc time)

julia> m1000 = f(1000)
Model{typeof(f), (:N, :σ), (), (), Tuple{Int64, Float64}, Tuple{}, DefaultContext}(f, (N = 1000, σ = 0.1), NamedTuple(), DefaultContext())

julia> @be DynamicPPL.predict(m1000, fchain; include_all=false)
Benchmark: 3 samples with 1 evaluation
        27.579 ms (202050 allocs: 63.224 MiB, 20.14% gc time)
        46.832 ms (202050 allocs: 63.224 MiB, 52.75% gc time)
        163.931 ms (202050 allocs: 63.224 MiB, 86.04% gc time)

julia> @be DynamicPPL.predict(m1000, mchain; include_all=false)
Benchmark: 1 sample with 1 evaluation
        2.648 s (25296094 allocs: 826.041 MiB, 9.72% gc time, without a warmup)