Skip to content

feat: backend switching for Mooncake #768

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 31 commits into from
Jun 18, 2025
Merged

Conversation

AstitvaAggarwal
Copy link
Contributor

@AstitvaAggarwal AstitvaAggarwal commented Apr 1, 2025

Define Mooncake.rrule!! for DI.DifferentiateWith.

@AstitvaAggarwal AstitvaAggarwal requested a review from gdalle as a code owner April 1, 2025 12:30
Copy link

codecov bot commented Apr 1, 2025

Codecov Report

Attention: Patch coverage is 96.82540% with 2 lines in your changes missing coverage. Please review.

Project coverage is 97.87%. Comparing base (eeb4c86) to head (1e85f17).

Files with missing lines Patch % Lines
...ationInterface/test/Back/DifferentiateWith/test.jl 88.23% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #768      +/-   ##
==========================================
- Coverage   97.88%   97.87%   -0.01%     
==========================================
  Files         128      129       +1     
  Lines        7694     7754      +60     
==========================================
+ Hits         7531     7589      +58     
- Misses        163      165       +2     
Flag Coverage Δ
DI 98.72% <96.82%> (-0.03%) ⬇️
DIT 95.61% <ø> (ø)

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.

@gdalle gdalle marked this pull request as draft April 1, 2025 12:43
@AstitvaAggarwal
Copy link
Contributor Author

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

Copy link
Member

@gdalle gdalle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Contributor Author

AstitvaAggarwal commented Apr 9, 2025

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

@gdalle
Copy link
Member

gdalle commented Apr 10, 2025

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

Copy link
Member

@gdalle gdalle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

@yebai yebai left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good work, @AstitvaAggarwal. We are getting there. I have just a few more quality-related comments on the documents, type stability, etc.

Copy link
Contributor

@yebai yebai left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @AstitvaAggarwal for the nice work!

EDIT: I am happy with this PR, but still need @gdalle's blessing before merging.

@yebai yebai marked this pull request as ready for review June 9, 2025 21:50
Copy link
Member

@gdalle gdalle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very good work @AstitvaAggarwal, thank you so much for taking the time! My main remaining questions are about:

  • support for tuples: do we keep it, do we remove it? It seems a bit arbitrary to have just tuples in addition to the "officially" supported Number and AbstractArray, especially if these can only be tuples of numbers
  • testing code, which should not touch on so many Mooncake internals

To lighten the load on you, I'll make the suggested modifications myself, and you can tell me what you think after! Thanks again for your patience.

Comment on lines 3 to 5
# nested vectors (eg. [[1.0]]), Tuples (eg. ((1.0,),)) or similar (eg. [(1.0,)]) primal types are not supported by DI yet !
# This is because basis construction (DI.basis) does not have overloads for these types.
# For details, refer commented out test cases to see where the pullback creation fails.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The issue is that we're testing DifferentiateWith(f, substitute_backend) with substitute_backend = AutoFiniteDiff(), aka a forward-mode backend. I think it should work with DifferentiateWith(f, AutoEnzyme())?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolved by removing tuples for the time being

@@ -0,0 +1,268 @@
@is_primitive MinimalCtx Tuple{DI.DifferentiateWith,<:Union{Number,AbstractArray,Tuple}}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I find it a bit weird to have this union of Number, AbstractArray (two types which are theoretically supported for x inputs in DI) and then just Tuple (which is not officially part of the supported inputs). Why not also NamedTuple for instance? Is it better if we just say Any? Or restrict to Number and AbstractArray for the time being?

@gdalle gdalle changed the title Backend switching for Mooncake feat: backend switching for Mooncake Jun 13, 2025
@gdalle gdalle added the skipci Don't run CI tests and docs label Jun 18, 2025
@gdalle gdalle merged commit 8c234e6 into JuliaDiff:main Jun 18, 2025
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
skipci Don't run CI tests and docs
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants