Skip to content
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

Remove BijectorsEnzymeExt on 1.11.1 #337

Merged
merged 5 commits into from
Oct 29, 2024
Merged

Conversation

penelopeysm
Copy link
Member

@penelopeysm penelopeysm commented Oct 28, 2024

Description of old PR, no longer relevant

See #333 and the comment added to BijectorsEnzymeExt for context.

I can confirm that this doesn't error in all contexts (but I can't confirm that it has the desired effect, because we can't actually meaningfully test Enzyme on 1.11):

  • Precompilation: doesn't crash
  • Runtime using Enzyme; using Bijectors: doesn't crash
  • Runtime using Bijectors; using Enzyme: doesn't crash

I tested this locally on Julia 1.11.1 and [email protected]. I've also tested on Julia 1.10.6.

In theory, the main downside of this is:

Overall, though, I think this is a better fix than #333, because it does address the root problem of us implicitly depending on the load order of the two extensions. It just introduces other problems (of maintainability) in the process, but hopefully the comment is clear enough and mitigates this.

Closes #332 by simply removing the code on 1.11.1+.

@penelopeysm penelopeysm force-pushed the py/inline-import-rules branch 2 times, most recently from 30f7ff7 to 6b78489 Compare October 28, 2024 20:22
@penelopeysm penelopeysm marked this pull request as draft October 28, 2024 20:42
@yebai
Copy link
Member

yebai commented Oct 28, 2024

Feel free to disable these Enzyme rules and the code that depends on them.

We can reactivate these codes after JuliaLang/julia#56368 is merged and released.

@penelopeysm
Copy link
Member Author

penelopeysm commented Oct 28, 2024

Thanks for finding that @yebai! That's exciting.

I wonder what happens on 1.10 though because the breakage on 1.10 isn't caused by the dependency chain but rather due to an inability to simultaneously have something as a dep of a main package and a dep of an extension (in our case, ChainRulesCore is both a dep of Bijectors, and of BijectorsEnzymeExt). Will have to do some testing.

@penelopeysm penelopeysm force-pushed the py/inline-import-rules branch 2 times, most recently from 34bfa83 to 4ab2a17 Compare October 29, 2024 02:42
@penelopeysm
Copy link
Member Author

penelopeysm commented Oct 29, 2024

The 1.10 failure should be fixed now through a devious workaround of using Bijectors: ChainRulesCore rather than using ChainRulesCore itself.

@penelopeysm
Copy link
Member Author

penelopeysm commented Oct 29, 2024

@willtebbutt @mhauru, I'm optimistic about this one 😄

CI is still failing for 1.11, but that's no longer because of the extension precompilation issue; it's because it's trying to load Tapir, which naturally doesn't work on 1.11.

#338 shows what happens if we were to replace Tapir -> Mooncake; I've kept it as a separate PR to make the diffs clearer. On that PR all CI tests pass except for:

  • Enzyme on 1.11 (which we know is broken; we could disable it from running in CI too)
  • one Mooncake error (which I haven't looked into, but since I have your eyes on it here...)
  • one of the interface tests, which looks to me like a case of unlucky random number generation (which we can fix later by seeding the RNG).

If you like the looks of both, I suggest we could merge #338 into this, merge this into master, open issues for the other failing tests, then tag a new release, so that the precomp issue doesn't hold up other work.

@penelopeysm penelopeysm mentioned this pull request Oct 29, 2024
@penelopeysm penelopeysm marked this pull request as ready for review October 29, 2024 03:30
Copy link
Member

@mhauru mhauru left a comment

Choose a reason for hiding this comment

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

Assuming the unlucky RNG test passes (I'm rerunning it) I'd be happy to merge #338 and this. Marking Enzyme v1 CI to not run would be good too.

I think we should revisit this once hopefully Julia v1.11.2 fixes the problem for good, the right way (JuliaLang/julia#56368), and maybe by then Enzyme will be working with v1.11 too. Might end up keeping the code introduced in this, but with a version bound of ==v1.11.1, or might even remove it to simplify and raise an error saying "if you want to use Enzyme with Bijectors, you need some other, any other, version". Regardless, this seems like definitely an improvement over the current state, and good enough for now.

@penelopeysm penelopeysm changed the title Directly inline expansion of @import_rrule and @import_frule into BijectorsEnzymeExt Remove BijectorsEnzymeExt on 1.11.1 Oct 29, 2024
@penelopeysm
Copy link
Member Author

The code is removed now (and replaced with a warning). Same deal as before – we should check CI on #338, merge that into this, and then merge this into master and release.

Copy link
Member

@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 @penelopeysm!

* Tapir -> Mooncake

* Bump minor version

* Mark Mooncake test as broken

* Remove BijectorsEnzymeExt on 1.11.1+

* Increase tolerance on `ordered` test
@penelopeysm penelopeysm merged commit e0f04fc into master Oct 29, 2024
21 of 25 checks passed
@penelopeysm penelopeysm deleted the py/inline-import-rules branch October 29, 2024 16:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Bijectors 0.13.18 not working with Enzyme on Julia 1.11.1
3 participants