Tangent Interface Docs

[willtebbutt]

@Jollywatt per our discussion elsewhere, here is a stab at aggregating all of the information that I have about the tangent interface into a single location. When you get a minute, could you let me know if this is the kind of thing that you need to do what you wanted to do, or if there's additional info that you'll need?

[codecov]

Codecov Report

Attention: Patch coverage is 50.00000% with 13 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/test_utils.jl	54.16%	11 Missing ⚠️
src/tangents.jl	0.00%	2 Missing ⚠️

Files with missing lines	Coverage Δ
src/fwds_rvs_data.jl	20.64% <ø> (-76.03%)	⬇️
src/tangents.jl	10.24% <0.00%> (-74.99%)	⬇️
src/test_utils.jl	30.03% <54.16%> (-63.07%)	⬇️

... and 35 files with indirect coverage changes

[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 │     84.8 │     1.0 │        5.61 │    8.21 │
│                  _sum_1000 │      6.6 │  1380.0 │        34.4 │    1.09 │
│               sum_sin_1000 │     2.29 │    1.68 │        10.8 │    1.96 │
│              _sum_sin_1000 │     2.71 │   255.0 │        13.6 │    2.44 │
│                   kron_sum │     62.0 │    3.55 │       232.0 │    8.44 │
│              kron_view_sum │     21.3 │    3.32 │        80.8 │    34.0 │
│      naive_map_sin_cos_exp │      2.6 │ missing │        7.63 │    2.35 │
│            map_sin_cos_exp │     2.72 │    1.44 │        6.01 │    2.94 │
│      broadcast_sin_cos_exp │     2.59 │    2.23 │        1.46 │    2.24 │
│                 simple_mlp │     4.78 │    3.25 │        6.45 │     3.3 │
│                     gp_lml │     13.2 │    6.55 │     missing │     6.2 │
│ turing_broadcast_benchmark │     3.22 │ missing │        25.2 │ missing │
│         large_single_block │     4.42 │  4380.0 │        30.6 │    2.18 │
└────────────────────────────┴──────────┴─────────┴─────────────┴─────────┘

[sunxd3]

I haven't tried to implement a tangent type and use these functions, so couldn't verify claims like If all the tests in these functions pass, then you have satisfied the interface. All the comments are only cosmetic.

But this is clearly a really good clarification!

[sunxd3]

maybe miss ref for rdata_type?

[willtebbutt]

Ah, yes. For some reason, it didn't like me doing that. I'm not 100% sure why.

[Jollywatt]

@willtebbutt I've had a look at these docs and tried to implement the interface for a custom type. I was almost able to, but in particular I had to add methods for

• _get_fdata_field
• tangent (binary method)

in addition to those in the interface. I also ran into problems with a generated function, which I didn't find a way around.

I've put a MWE in this notebook: https://jollywatt.github.io/notes/mooncake-self-tangents

[willtebbutt]

Thanks -- I'll take a look at your notebook!

[yebai]

@Jollywatt @willtebbutt I suggest that you take a look at #428 (comment) as part of "integration testing" this PR's completeness and accuracy. Defining a macro that can help users extend tangent types for their Julia types will be nice.