Handle Unions / UnionAlls of Tuples in zero_rdata

[willtebbutt]

Resolves #283

[willtebbutt]

@penelopeysm when you get a minute, could you confirm that this resolves your issue #283 ? I just missed a case in some of my internal logic, and your example appears to be the first time that it's been encountered in practice.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Files with missing lines	Coverage Δ
src/fwds_rvs_data.jl	96.48% <100.00%> (+0.01%)	⬆️

[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 │ Mooncake │  Zygote │ ReverseDiff │  Enzyme │
│                     String │   String │  String │      String │  String │
├────────────────────────────┼──────────┼─────────┼─────────────┼─────────┤
│                   sum_1000 │    112.0 │     1.0 │        5.51 │     1.9 │
│                  _sum_1000 │     7.74 │  1380.0 │        34.8 │   0.106 │
│               sum_sin_1000 │     2.54 │     1.6 │        11.4 │    1.01 │
│              _sum_sin_1000 │     2.88 │   297.0 │        16.4 │    1.38 │
│                   kron_sum │     60.5 │    10.2 │       218.0 │    11.5 │
│              kron_view_sum │     28.8 │    4.07 │        77.8 │    4.19 │
│      naive_map_sin_cos_exp │     3.18 │ missing │        8.91 │    2.82 │
│            map_sin_cos_exp │     4.71 │    1.71 │        7.45 │    3.45 │
│      broadcast_sin_cos_exp │     4.36 │    2.56 │        1.68 │    2.84 │
│                 simple_mlp │     8.56 │    3.28 │        12.7 │    3.32 │
│                     gp_lml │     16.9 │    6.68 │     missing │ missing │
│ turing_broadcast_benchmark │     7.21 │ missing │        35.9 │ missing │
│         large_single_block │     3.89 │  4100.0 │        31.0 │    2.23 │
└────────────────────────────┴──────────┴─────────┴─────────────┴─────────┘

[penelopeysm]

Thanks for fixing this!! Yes it does fix all the issues with constrained distributions in the Turing test suite too – https://github.com/TuringLang/Turing.jl/actions/runs/11214726985/job/31170411370 🙂

There is another error which I didn't notice before because it was in the same test group as the other errors 😬 will open another issue but I think it's safe to say that we're good on #283!

[willtebbutt]

Thanks -- that looks like a weird missing rule problem (weird because I thought I had all of the Cholesky stuff covered, but I guess not). I'll merge this + close #283 . Thanks for checking!