Further ChainRulesCore.rrule Integration

[willtebbutt]

At present the extent to which ChainRulesCore.rrules can be (straightforwardly) used in this package is quite limited. The purpose of this PR is to:

• make it easier to incorporate more interesting rules, thus closing Improving ChainRules rule usage #191 ,
• include some failing tests to ensure that we get clear errors when conversations between Tapir and ChainRules tangent types goes wrong,
• explain the limitations of what can be done,
• provide guidance to the situations in which you should consider importing an rrule from ChainRules, and
• import rules from NNlib.jl and test that everything works NNlib Support #171 .
• reduce quantity of code duplication between rrule wrapper with / without kwargs.
• improve macro for use with kwargs rrules.
• improve macro with where terms so that diagonal dispatch can be done.
• improve chain rules interop documentation.
• add + test rules associated with LuxLib.jl
• integration test Lux.jl

The reason for incorporating the rules from NNlib.jl in this PR, rather than a subsequent one, is that the main place we wish to import rules from in the near term is NNlib.jl. Consequently, doing that in this PR will give us a good sense of whether the new functionality we introduce here is sufficient.

edit: as a happy consequence of thinking more carefully about how we incorporate functionality from ChainRules, I've removed a lot of code from the the @from_rrule macro, and put it in a regular Julia function. This makes this bit of the codebase much more straightforward to understand.

[codecov]

Codecov Report

Attention: Patch coverage is 98.55072% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/chain_rules_interop.jl	98.46%	1 Missing ⚠️

Files with missing lines	Coverage Δ
src/Mooncake.jl	100.00% <ø> (ø)
src/rrules/blas.jl	98.08% <100.00%> (+0.02%)	⬆️
src/chain_rules_interop.jl	98.46% <98.46%> (ø)

[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 │    105.0 │     1.0 │        5.61 │     2.0 │
│                  _sum_1000 │     7.76 │  1360.0 │        32.9 │   0.106 │
│               sum_sin_1000 │     2.56 │     1.6 │        10.7 │    1.03 │
│              _sum_sin_1000 │     3.17 │   322.0 │        17.7 │     1.5 │
│                   kron_sum │     59.7 │    9.96 │       212.0 │    10.1 │
│              kron_view_sum │     80.7 │    11.1 │       224.0 │    10.0 │
│      naive_map_sin_cos_exp │     3.17 │ missing │        8.83 │    2.78 │
│            map_sin_cos_exp │     4.64 │     1.7 │        7.52 │     3.4 │
│      broadcast_sin_cos_exp │     4.36 │    2.61 │         1.7 │    2.86 │
│                 simple_mlp │     5.28 │    3.24 │        7.09 │    3.27 │
│                     gp_lml │     5.57 │    2.21 │     missing │ missing │
│ turing_broadcast_benchmark │     7.12 │ missing │        31.4 │ missing │
│         large_single_block │      3.9 │  4440.0 │        30.3 │    2.24 │
└────────────────────────────┴──────────┴─────────┴─────────────┴─────────┘