Add DynamicPPL extension to allow rule-reusing with setmodel

[sunxd3]

TuringLang/DynamicPPL.jl#640

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/TapirDynamicPPLExt.jl	0.00%	6 Missing ⚠️

Files with missing lines	Coverage Δ
ext/TapirDynamicPPLExt.jl	0.00% <0.00%> (ø)

[github-actions]

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬────────┬─────────┬─────────────┬─────────┐
│                      Label │  Tapir │  Zygote │ ReverseDiff │  Enzyme │
│                     String │ String │  String │      String │  String │
├────────────────────────────┼────────┼─────────┼─────────────┼─────────┤
│                   sum_1000 │  115.0 │    0.79 │        4.81 │    1.71 │
│                  _sum_1000 │   7.93 │  1430.0 │        32.3 │  0.0841 │
│               sum_sin_1000 │   2.95 │    1.61 │        10.8 │    1.01 │
│              _sum_sin_1000 │   3.59 │   324.0 │        16.9 │    1.54 │
│                   kron_sum │   85.1 │    3.51 │       219.0 │    8.16 │
│              kron_view_sum │   97.2 │    11.4 │       238.0 │    9.01 │
│      naive_map_sin_cos_exp │   4.16 │ missing │        9.08 │    2.82 │
│            map_sin_cos_exp │   4.74 │    1.81 │        7.61 │    3.45 │
│      broadcast_sin_cos_exp │   4.62 │    2.66 │        1.66 │    2.83 │
│                 simple_mlp │    8.7 │    3.09 │        10.9 │    3.12 │
│                     gp_lml │   15.6 │    4.25 │     missing │ missing │
│ turing_broadcast_benchmark │   8.92 │ missing │        26.4 │ missing │
└────────────────────────────┴────────┴─────────┴─────────────┴─────────┘

[willtebbutt]

I'm not entirely sure that this solves the Gibbs sampling problem. To replicate, try something like

using ADTypes, Random, Tapir, Turing

@model function gdemo()
    s ~ InverseGamma(2, 3)
    m ~ Normal(0, sqrt(s))
    1.5 ~ Normal(m, sqrt(s))
    2.0 ~ Normal(m, sqrt(s))
    return s, m
end

rng = Xoshiro(123);
alg = Gibbs((PG(2_000, :s), 1), (HMCDA(500, 0.8, 0.25, :m; init_ϵ=0.05, adtype=ADTypes.AutoTapir(false)), 5_000));
sample(rng, gdemo(), alg, 5, discard_initial=2);

Also, if you add some kind of print statement in Tapir.build_rrule to display the sig_or_mi argument, you'll be able to visually track when it is that Tapir.jl is deriving a new rule.

This example highlights that it is indeed the Gibbs sampling aspect of this that's taking the time, because there are 5000 iterations of HMCDA at each step of Gibbs sampling, which do not appear to take much time at all (Particle Gibbs is quite expensive though, so I'm only doing a single iteration of that).

[sunxd3]

This PR doesn't fully resolve the rule re-deriving problem (ref TuringLang/Turing.jl#2289 (comment)).

Efforts like #237 seems to be the way for the future.