Backend switching for Mooncake

[AstitvaAggarwal]

Define Mooncake.rrule!! for DI.DifferentiateWith.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.05%. Comparing base (4b85c4f) to head (233c312).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #768      +/-   ##
==========================================
- Coverage   97.93%   97.05%   -0.89%     
==========================================
  Files         129      127       -2     
  Lines        7458     7466       +8     
==========================================
- Hits         7304     7246      -58     
- Misses        154      220      +66     

Flag	Coverage Δ
DI	97.86% <100.00%> (-1.16%)	⬇️
DIT	95.01% <ø> (-0.20%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

🚀 New features to boost your workflow:

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

[AstitvaAggarwal]

removed the code that piggybacks off the Chainrules wrapper. This is specifically now a Mooncake generic rule which handles backend switching.

[gdalle]

Thanks for this first draft!
I think there are some changes necessary, and most importantly you need to test it, first locally and then during CI (try not to run CI before having tested your changes locally, the process is very expensive since it tests a dozen different backends for like half an hour each).
For the testing, start with manual tests, and then once your code works you can add AutoMooncake() to this line

[AstitvaAggarwal]

sorry i got preoccupied with some other work, hence the incomplete PR. This would be on route now.

[gdalle]

Please keep in mind that every commit costs around 6 hours of CI budget. I suggest you make as many modifications as possible locally and add tests first before pushing

[gdalle]

Thanks, we're getting closer!
Unfortunately I think my existing tests are not enough to capture everything that can go wrong in a Mooncake rule. Perhaps the Mooncake test utilities should be brought in, or more sophisticated tests should be written.

[gdalle]

I'm not sure about this, to me it seems like we're forgetting to carry around the FData of y.

[AstitvaAggarwal]

im not sure i understand what you mean by this. what ive tried here is instead of creating the output Codual later, we create it early so that in cases of memory address function return types, the pullback functions can have access to the FData to create the correct adjoints.

[gdalle]

What if there are derivatives inside dw?

[AstitvaAggarwal]

im not sure i understand. dw would be a function and if derivatives are used inside i think that would be handled by the substitute backend?

[gdalle]

We're not sure that f has no RData

[AstitvaAggarwal]

I think all functions would have no RData

[gdalle]

Not all callable objects are functions though.

julia> struct Multiplier
           a::Float64
       end

julia> (m::Multiplier)(x) = m.a * x

julia> m = Multiplier(2)
Multiplier(2.0)

julia> m(3)
6.0

julia> m isa Function
false

[gdalle]

And I'm not even convinced the current Mooncake behavior is correct, see compintell/Mooncake.jl#557

[AstitvaAggarwal]

Nice catch, Mooncake treats all callable objects as a function, giving NoRData(). The issue you have opened in mooncake, i think would need to be figured out independent of DI (atleast this PR) as of now (comes in when Mooncake is a substitute backend as well).