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 89.87342% with 16 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/test_utils.jl	89.18%	8 Missing ⚠️
src/tangents.jl	86.95%	3 Missing ⚠️
src/rrules/iddict.jl	80.00%	2 Missing ⚠️
src/rrules/tasks.jl	75.00%	2 Missing ⚠️
src/rrules/memory.jl	95.45%	1 Missing ⚠️

Files with missing lines	Coverage Δ
ext/MooncakeCUDAExt.jl	93.33% <100.00%> (ø)
src/fwds_rvs_data.jl	86.40% <ø> (ø)
src/rrules/array_legacy.jl	62.79% <100.00%> (ø)
src/rrules/function_wrappers.jl	95.95% <100.00%> (ø)
src/rrules/twice_precision.jl	96.15% <100.00%> (ø)
src/rrules/memory.jl	95.61% <95.45%> (ø)
src/rrules/iddict.jl	95.55% <80.00%> (ø)
src/rrules/tasks.jl	79.68% <75.00%> (ø)
src/tangents.jl	86.34% <86.95%> (+0.24%)	⬆️
src/test_utils.jl	91.78% <89.18%> (-1.76%)	⬇️

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[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 │   Primal │ Mooncake │  Zygote │ ReverseDiff │  Enzyme │
│                     String │   String │   String │  String │      String │  String │
├────────────────────────────┼──────────┼──────────┼─────────┼─────────────┼─────────┤
│                   sum_1000 │ 100.0 ns │      1.8 │     1.0 │        5.61 │    8.21 │
│                  _sum_1000 │ 941.0 ns │     6.67 │  1620.0 │        33.1 │    1.09 │
│               sum_sin_1000 │  6.54 μs │      2.2 │    1.69 │        10.6 │    1.97 │
│              _sum_sin_1000 │  5.25 μs │     2.54 │   288.0 │        13.2 │    2.43 │
│                   kron_sum │ 361.0 μs │     35.5 │    12.3 │       199.0 │    7.69 │
│              kron_view_sum │ 361.0 μs │     38.9 │    9.87 │       215.0 │    91.4 │
│      naive_map_sin_cos_exp │  2.16 μs │     2.17 │ missing │        7.07 │    2.33 │
│            map_sin_cos_exp │  2.14 μs │     2.38 │    1.57 │        6.07 │    2.88 │
│      broadcast_sin_cos_exp │  2.23 μs │     2.26 │    2.38 │        1.47 │    2.26 │
│                 simple_mlp │ 163.0 μs │     7.06 │     1.7 │        12.7 │    3.84 │
│                     gp_lml │ 194.0 μs │     10.7 │    4.39 │     missing │    5.48 │
│ turing_broadcast_benchmark │  1.97 ms │     3.52 │ missing │        24.6 │ missing │
│         large_single_block │ 390.0 ns │     4.39 │  4410.0 │        32.9 │    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!

[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.

[willtebbutt]

@Jollywatt sorry for taking so long to get this fixed. I believe I've addressed the issues that you mentioned above.

In particular, I've added the binary version of tangent to the docstring for the tangent splitting tests, and stated in the docstring from the basic rules tests that you need to support _new_ / getfield / setfield! etc (which is where the _get_fdata_field thing you mentioned is coming from).

Could you let me know if you think this is sufficient? If so, I'll merge this.

[Jollywatt]

@willtebbutt I'm just having another shot at implementing the tangent interface following the docstring for Mooncake.TestUtils.test_tangent_interface.

Last time, test_tangent_interface tested whether methods for the "user-facing" methods (e.g., tangent_type, zero_tangent, …) worked as expected. However, now it directly tests internal functions (e.g., zero_tangent_internal) so doesn't test methods I define on the user-facing functions.

Am I meant to use Mooncake.TestUtils.test_tangent_interface to test methods on user-facing functions like I did last time in https://jollywatt.github.io/notes/mooncake-self-tangents, or has that changed?

[willtebbutt]

Ah, yes, you're definitely meant to test on the <function>_internal functions, but I've not updated this in the various docstrings 😬 . I forgot that I would need to reflect this change in them. I'll update them and let you know once that's done.

[willtebbutt]

@Jollywatt I think that the docstring should now be accurate.