Add easy way to use Mooncake.jl for gradients

[mcabbott]

This is parallel to FluxML/Flux.jl#2471 for Enzyme. See compintell/Mooncake.jl#361 for some benchmarks (without using this).

Questions for @willtebbutt are:

• Do you hate this? Should Mooncake own the struct? It's not far from CoDual but (1) we want a one-arg method, and (2) Flux wants nicer printing, and (3) that is marked private & discouraged?
• How do I implement Const, to take gradient with respect to only some arguments? DI seems to just compute them all & discard. CoDual(x, NoTangent()) does not seem to work.

Re the first point, rather than having one struct per package, another option might be to have one universal thing:

GradWrap(Enzyme, model)    # returns something that behaves like Duplicated(model)
GradWrap(Mooncake, model)  # instead of Moonduo
GradWrap(Reactant, model)  # for Reactant.jl, would store compiled functions not just gradients, mutable

Flux.gradient(loss, ::GradWrap, ...)  # uses selected package

[codecov]

Codecov Report

Attention: Patch coverage is 45.00000% with 33 lines in your changes missing coverage. Please review.

Project coverage is 73.55%. Comparing base (21693da) to head (139ff3e).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
ext/FluxMooncakeExt.jl	51.02%	24 Missing ⚠️
src/mooncake.jl	18.18%	9 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master      #26      +/-   ##
==========================================
- Coverage   79.92%   73.55%   -6.37%     
==========================================
  Files           7        9       +2     
  Lines         269      329      +60     
==========================================
+ Hits          215      242      +27     
- Misses         54       87      +33     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

---

🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests

[willtebbutt]

This is cool. Thanks for engaging with Mooncake on this!

In answer to your questions:

1. I do not hate this at all. It's not like it's lots of code duplication, and you having ownership over the type seems very reasonable to me.
2. Sadly, the correct thing to do is currently to do what DI currently does, and simply ignore the terms you do not want. I hope this will change at some point, but I have no idea when that will happen.

edit: also, you're correct, CoDual(x, NoTangent()) won't work in general. Mooncake's semantics are: if you see a CoDual(x, dx), it must be true that either

1. typeof(dx) == tangent_type(typeof(x)), or
2. typeof(dx) == fdata_type(tangent_type(typeof(x))).

Anything else is incorrect, and will cause you problems. Which one of these you should find depends on context. In all of the functionality that you are using, it should always be 1.

[willtebbutt]

To the best of my knowledge, this shouldn't ever be a case that you see.

[mcabbott]

Ok, I didn't try very hard! Do you know if there are other types which can occur, besides those handled here?

[willtebbutt]

Do you place any constraints on the types of the things that you take gradients w.r.t.?

I ask because there's nothing to stop people adding a new type, and declaring a non-standard tangent type for it (i.e. while some subtype of Tangent is the thing returned by tangent_type by default for structs, there's nothing to stop people making tangent_type return something else for a type that they own). So in principle you could see literally anything. In practice, assuming that you're working with Arrays, and structs / mutable structs / Tuples / NamedTuples of Arrays, I think you should be fine.

My honest advice would be to do what you're doing at the minute. i.e. check that it works for models that you care about, and ensure that there's a good error message so that a user knows where to ask for help if they encounter something you weren't expecting.

[mcabbott]

Ok, sounds good!

No constraints, but in reality it's going to be arrays & all kinds of structs.

[willtebbutt]

Is there any chance at all that f will contain trainable parameters, or does Flux insist that you not do that?

[mcabbott]

Yea Flux.gradient(f, x, y) has length 2, alla Zygote. I agree that's not the maximally flexible thing, and very occasionally you end up with gradient(|>, x, f)... but in real use it seems like never.

[willtebbutt]

Cool

[mcabbott]

Thanks for having a look!

Ok, it is easy of course to hide undesired inputs. That's all that's possible with Zygote too.

I could check but do you know whether gradient(x -> f(x, y), x) always does the work to compute dy, or does it notice that this is global & propagate its constant-ness along? Zygote always computes. If Mooncake does not, then we (and DI) could transform gradient(f, CoDual(x, ...), Const(y)) to that.

[willtebbutt]

It's always computed I'm afraid -- since x -> f(x, y) is just a closure, and therefore has a y field, Mooncake treats it like any other struct, and computes the gradient w.r.t. the parameters it closes over.