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

Disable precompilation of BijectorsEnzymeExt on Julia 1.11+ #333

Closed
wants to merge 4 commits into from

Conversation

penelopeysm
Copy link
Member

@penelopeysm penelopeysm commented Oct 24, 2024

In the following, B = Bijectors, E = Enzyme, and C = ChainRulesCore.

We have BEExt in this repository, and Enzyme.jl defines ECExt, and B depends on C. The problem is that BEExt depends on a method that is defined in ECExt, and therefore requires ECExt to have been loaded before BEExt.

In Julia 1.10.5 and 1.11.0, this requirement was somehow always satisfied. (I don't know why.)

However, JuliaLang/julia#55589 shook this up and now, there's no guarantee that ECExt will be loaded first. Indeed, it turns out that in Julia 1.11.1, BEExt is always loaded first. The PR is slated to be backported to 1.10.6 as well. (It hasn't been backported.) I suppose that we cannot really complain about the PR because the dependency graph in itself doesn't specify that one should be loaded before the other (the order of loading depends only on how the dependency graph is traversed), and we probably should not be relying on that implicit behaviour. Although the discussion in JuliaLang/julia#55589 suggests that we might one day be able to explicitly specify dependencies between extensions (i.e. adding an extra edge to the dependency graph), this isn't in place yet.

Anyway, the result of this is that precompilation fails when adding Bijectors and Enzyme to an environment, as reported in #332.


This PR uses a workaround described in JuliaLang/julia#56204 to disable precompilation for BEExt, so that we can at least add Bijectors and Enzyme to the same environment.

It doesn't fully solve the issue at runtime, though. At runtime, we can – to some extent – control the order of loading. If we run using E before using B, then ECExt gets loaded first, and all is good. But if we run using B first, then BEExt is loaded, and we get an error.

I don't know how to fully solve the runtime issue, but this PR will at least solve the precomp issue.


Closes #332 (in the sense that it fixes the immediate problem of precompilation failing; the root cause of this issue is really the underspecified dependency graph, but there isn't a fix for that)

@penelopeysm penelopeysm force-pushed the py/fix-bijectors-enzyme-ext branch 6 times, most recently from 5086c31 to 95f8888 Compare October 24, 2024 01:35
@penelopeysm penelopeysm force-pushed the py/fix-bijectors-enzyme-ext branch 2 times, most recently from 6a42054 to a962770 Compare October 24, 2024 02:03
@penelopeysm
Copy link
Member Author

Note that to get CI to pass, we probably need to update Tapir to Mooncake. I started some work on this here https://github.com/TuringLang/Bijectors.jl/tree/py/tapir-to-mooncake but would prefer to keep that to a separate PR.

@wsmoses
Copy link
Collaborator

wsmoses commented Oct 24, 2024

This doesn’t really fix things so much as cause more breakage for users thereof

you can specify a package to have multiple dependencies to ensure loading. If you mark the enzyme ext as requiring both Enzyme and CR, does that resolve. And/or actually loading CR in the extension

@penelopeysm
Copy link
Member Author

This doesn’t really fix things so much as cause more breakage for users thereof

Can you explain?

If you mark the enzyme ext as requiring both Enzyme and CR, does that resolve. And/or actually loading CR in the extension

Neither of those work, unfortunately!

@wsmoses
Copy link
Collaborator

wsmoses commented Oct 24, 2024

yeah so for enzyme users here, now its much more broken =/

What does CI say when you do the corresponding fixes above?

@penelopeysm
Copy link
Member Author

penelopeysm commented Oct 24, 2024

yeah so for enzyme users here, now its much more broken =/

I'm very unfamiliar with Enzyme and the AD ecosystem, so I'm afraid you might have to spell this out more clearly for me.

How is it broken? Or to be more precise, it is definitely still somewhat broken (I wrote as much in my original comment), but how is it much more broken than the current version?

What does CI say when you do the corresponding fixes above?

Precompilation fails in all cases: #334 #335 #336 and it's reproducible locally too.

@penelopeysm
Copy link
Member Author

Disabled fast failing on CI so that we can see exactly which tests error. They're fairly cheap to run anyway.

@mhauru
Copy link
Member

mhauru commented Oct 25, 2024

Depending on import order in one's code I think is pretty horrendous, and we should make a proper fix a priority. However, this seems like an improvement to me, at least allowing environments with both packages to precompile. Would like to hear why @wsmoses considers the partial fix worse than no fix at all though, before approving.

@penelopeysm
Copy link
Member Author

Closing in favour of #337, which imo is a better fix

@yebai yebai deleted the py/fix-bijectors-enzyme-ext branch October 29, 2024 17:10
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