Explicit Enzyme rules on Enzyme 0.13 #350
Conversation
* Bump minimum Julia version to 1.10 * Use 'min' in CI * Bump versions of GHA (julia-actions needs to be v2 for 'min') * Update Project.toml --------- Co-authored-by: Hong Ge <[email protected]>
|
The tests of the forward-mode rule pass on Julia 1.10 but fail on Julia 1.11 for all- |
|
cc @wsmoses |
|
huh, yeah looks like we've never had anyone define a reverse-mode rule with a batched active [ie non duplicated return]. I've made it a beter error message in Enzyme. Probably just disable that test. And open a mwe for the 1.11 issue? |
|
Here's an issue for the error that only occurs on Julia v1.11: EnzymeAD/Enzyme.jl#2121 This might not be the root issue here though, like @devmotion says it seems like the rule defined here might just not get picked up? |
|
@mhauru can you open an issue for the fact the rule isn't getting picked up? |
I tried to put together a small MWE for the issue last week but of course it worked without any problem 😅 After disabling the problematic tests, the EnzymeTestUtils tests of the rules pass. However, it seems Enzyme errors (in a quite nasty way) on Julia 1.11 (other Julia versions are fine) when running the the closely related tests in https://github.com/TuringLang/Bijectors.jl/blob/master/test/ad/flows.jl. Maybe it's also due to the rule not being picked up for some activities? |
I looked at that issue it it appears to be a bug in Julia itself. If you can open an issue for the lacking rule [maybe even start by taking bijectors itself and removing code as much as possible?], that would be helpful so we can make sure the rule is triggered for 1.11 In any case this PR lgtm and we can always follow it up |
test/ad/utils.jl
Outdated
| if !(:EnzymeForwardCrash in broken) | ||
| if forward_broken | ||
| @test_broken( | ||
| collect(et, Enzyme.gradient(Enzyme.Forward, f, x)[1]) ≈ finitediff, |
There was a problem hiding this comment.
this collect shouldn't be required anymore
test/ad/utils.jl
Outdated
| ) | ||
| else | ||
| @test( | ||
| collect(et, Enzyme.gradient(Enzyme.Forward, f, x)[1]) ≈ finitediff, |
It seems we might be able to merge the PR, @yebai? The rules themselves and the update to Enzyme 0.13 seems fine in principle according to @wsmoses (thanks a lot for your help on the way!), the remaining test Enzyme test failure on 1.11 seems to be an upstream issue. |
|
Thanks, @devmotion and @wsmoses -- it sounds good to merge this first. |
I rebased #349 on #341 and adapted the PR to Enzyme 0.13.