-
Notifications
You must be signed in to change notification settings - Fork 155
Add Mooncake integration #1342
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
Closed
Closed
Add Mooncake integration #1342
Changes from all commits
Commits
Show all changes
13 commits
Select commit
Hold shift + click to select a range
393a9f2
Add Mooncake integration
gdalle 6e31ec7
Restructure tests
gdalle 7a38354
Fix Ci
gdalle 2210dc9
Rename job
gdalle f9821cf
Path in quotes
gdalle 8d1b4c4
Fix path
gdalle 91e0cf4
Add caching
gdalle e1b3dd9
Remove Pkg from test deps
gdalle 4952e68
Fix
gdalle 099050d
Env in matrix
gdalle 880e669
Rename test
gdalle af7d1a8
Rerename
gdalle 5b73a34
Reactivate other tests
gdalle File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -7,3 +7,4 @@ | |
| /benchmark/*.json | ||
| /benchmark/Manifest.toml | ||
| /docs/Manifest.toml | ||
| /test/downstream/**/Manifest.toml | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,13 @@ | ||
| module StaticArraysMooncakeExt | ||
|
|
||
| using StaticArrays: SArray,MArray | ||
| using Mooncake: Mooncake | ||
|
|
||
| @static if isdefined(Mooncake, :FriendlyTangentCache) # checks Mooncake >= v0.5.25 | ||
| # see https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/998 | ||
| function Mooncake.friendly_tangent_cache(x::Union{SArray,MArray}) | ||
| return Mooncake.FriendlyTangentCache{Mooncake.AsPrimal}(copy(x)) | ||
| end | ||
| end | ||
|
|
||
| end | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,10 @@ | ||
| [deps] | ||
| Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" | ||
| StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" | ||
| Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" | ||
|
|
||
| [sources] | ||
| StaticArrays = {path = "../../../.."} | ||
|
|
||
| [compat] | ||
| Mooncake = ">=0.5.25" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,10 @@ | ||
| [deps] | ||
| Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" | ||
| StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" | ||
| Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" | ||
|
|
||
| [sources] | ||
| StaticArrays = {path = "../../../.."} | ||
|
|
||
| [compat] | ||
| Mooncake = "<0.5.25" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| using StaticArrays | ||
| using Mooncake | ||
| using Test | ||
|
|
||
| @testset verbose=true "Mooncake integration" begin | ||
| f(x) = sum(abs2, x) | ||
| config = Mooncake.Config(; friendly_tangents=true) | ||
| @testset "$(typeof(x))" for x in [ | ||
| SVector(1.0, 2.0), | ||
| MVector(1.0, 2.0) , | ||
| SMatrix{2,2}(1.0, 2.0, 3.0, 4.0), | ||
| MMatrix{2,2}(1.0, 2.0, 3.0, 4.0) | ||
| ] | ||
| cache = prepare_gradient_cache(f, zero(x); config) | ||
| val, grads = value_and_gradient!!(cache, f, x) | ||
| g = grads[2] | ||
| @test g isa typeof(x) | ||
| @test g == 2x | ||
| end | ||
| end |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| using Pkg | ||
| package = ENV["STATICARRAYS_DOWNSTREAM_TEST_PACKAGE"]; | ||
| env = ENV["STATICARRAYS_DOWNSTREAM_TEST_ENV"]; | ||
| Pkg.activate(joinpath(@__DIR__, package, env)) | ||
| Pkg.develop(PackageSpec(path=joinpath(@__DIR__, "..", ".."))) | ||
| include(joinpath(package, "test.jl")) |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this actually work with
SArraywhere elements areSymmetricSArrays?What exactly is Mooncake doing to the returned value to get a mutable type?
copyonSArrayreturns an immutable again.Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Probably not, one might need some kind of deep copy. Or we could restrict it to
Union{SArray{<:IEEEFloat}, MArray{<:IEEEFloat}}to be safe, albeit incomplete.Contrary to what the name "cache" suggests, I don't think the output of
friendly_tangent_cachehas to be mutable. From what I understand reading Claude's comments in chalk-lab/Mooncake.jl#1103, it seems that this function is meant to output a kind of template for reconstructing the gradient type we want.To clarify, I had nothing to do with the breaking changes in Mooncake v0.5.25, I actually disapprove of them and I don't fully understand them (especially since they were LLM-generated). I'm just trying to keep the bare minimum working (Mooncake's gradient returning a
SArrayin simple cases)There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't know enough about Mooncake to have an opinion yet. I just think a few more complex examples regarding how this is all supposed to work for nested types would be really helpful.
I did take a look at the Julia AD ecosystem a couple of years ago so I know a bit about how it works, and I still use it sometimes, though I don't see good areas to contribute to there sadly. Too many conflicting goals and points of view on how things should work.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think @yebai is the expert on this new Mooncake API. Perhaps it would be better as a Mooncake extension, since it is a very modest amount of code and StaticArrays is a very old and stable package?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
StaticArrays.jl already has a ChainRulesCore extension so adding Mooncake wouldn't be unprecedented. My main worry is that this friendly tangent thing doesn't seem particularly stable and well-tested, even in comparison with ChainRules. So it's unclear many iterations on the idea are still needed. I'd suggest figuring out more complex examples outside of StaticArrays.jl first, and make an extension when it's clear that the API works well across multiple nested array types from different packages.