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 (ef77d13) to head (36da036).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #768      +/-   ##
==========================================
- Coverage   97.93%   97.05%   -0.88%     
==========================================
  Files         128      126       -2     
  Lines        7693     7705      +12     
==========================================
- Hits         7534     7478      -56     
- Misses        159      227      +68     

Flag	Coverage Δ
DI	97.70% <100.00%> (-1.11%)	⬇️
DIT	95.33% <ø> (-0.29%)	⬇️

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]

In this file we need to add tests that are specific to Mooncake. Ideally I should have done that with the other backends too.
Basically, what we test now is the projection of the Mooncake rule you wrote onto the subset of stuff that DI cares about. But we should also check that the rule is correct from the Mooncake perspective. Probably the best tool for that is Mooncake.TestUtils.test_rule?

[AstitvaAggarwal]

yup, will check it out

[AstitvaAggarwal]

Okay so it fails for some of the general primal, functions. But for this PR its maybe okay? as DifferentiateWith is exclusive to DI, so the user is anyways limited to DI when using the Mooncake substitute backend.

[yebai]

it fails for some of the general primal, functions.

To clarify, these (primal) functions are permitted by the DI interface, right?

Assuming that is true, @gdalle, I think this is okay.

[yebai]

@AstitvaAggarwal is it possible to add these tests (excluding those not supported by DI) to this PR?