Fix tangent_type for Union{NoRData, RData{...}}

[yebai]

Summary

tangent_type(NoFData, R) previously only handled R <: Union{NoRData, IEEEFloat}, causing a missing-method error when R contained an RData{...} branch. This extends the dispatch to cover RData and adds _validate_rdata_union to catch invalid branches early.

Split from #1132 to make it easy to review.

Closes #1130.

CI Summary — GitHub Actions

Documentation Preview

Mooncake.jl documentation for PR #1133 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR1133/

Performance

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

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 180.0 ns │     1.51 │        1.61 │   0.778 │        3.56 │   1.39 │
│                  _sum_1000 │  1.07 μs │      6.8 │        1.02 │  4680.0 │        32.3 │   1.08 │
│               sum_sin_1000 │  7.42 μs │     2.58 │        1.18 │     1.7 │        9.06 │   1.85 │
│              _sum_sin_1000 │  5.11 μs │     3.57 │        2.47 │   296.0 │        13.3 │    2.8 │
│                   kron_sum │ 184.0 μs │     12.9 │        3.66 │    24.8 │       508.0 │   18.0 │
│              kron_view_sum │ 251.0 μs │     12.3 │        4.99 │    15.0 │       363.0 │   11.4 │
│      naive_map_sin_cos_exp │   2.3 μs │     2.69 │        1.48 │ missing │        7.19 │   2.07 │
│            map_sin_cos_exp │  2.26 μs │     3.12 │        1.53 │    2.43 │        5.95 │   2.67 │
│      broadcast_sin_cos_exp │  2.35 μs │     2.89 │        1.46 │    5.01 │        1.41 │   2.01 │
│                 simple_mlp │ 212.0 μs │     6.22 │        2.84 │    1.69 │        12.5 │   5.13 │
│                     gp_lml │ 400.0 μs │     5.91 │        1.87 │    6.57 │     missing │   2.63 │
│ turing_broadcast_benchmark │  2.05 ms │     5.05 │        3.32 │ missing │        32.6 │   1.79 │
│         large_single_block │ 450.0 ns │      5.7 │        1.98 │  3860.0 │        28.3 │   2.16 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[sunxd3]

thanks for initiating this!

[sunxd3]

I fear the tangent_type here can produce similar error as line 931, should we catch it?

[sunxd3]

it seems that Julia will still specialize the type instead of taking the Union type as hinted