Re-introducing Enzyme support on Julia 1.11

[penelopeysm]

This issue is intended to consolidate all the info about BijectorsEnzymeExt.

In the text that follows, to save the mental energy needed to parse long strings, B is Bijectors, E is Enzyme, and C is ChainRulesCore.

What happened with Bijectors v0.13.18?

BEExt uses the @import_frule and @import_rrule macros from Enzyme. While these macros are defined in Enzyme itself, under the hood they call the Enzyme._import_frule and Enzyme._import_rrule functions, which are defined inside ECExt.

Thus, if BEExt is loaded before ECExt, then a MethodError will be raised because the appropriate method has not yet been defined.

Up until Julia 1.11.0, there has been no need to worry about this, because Julia somehow always ensured that ECExt was loaded before BEExt. There is no logical reason for why one should be loaded before the other, but this behaviour has always been consistent. It is still true on Julia 1.10.6, as well.

However, Julia 1.11.1 (and specifically this PR) changed this 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, especially during precompilation. This meant that environments containing B and E would fail to precompile, as described in #332.

Several (partial) solutions to this problem were explored.

Solution 1: Disabling precompilation on 1.11.1+

See #333.

This workaround was suggested in this comment on the corresponding Julia issue. The result of this was that precompilation would succeed, and using E; using B at runtime would also succeed, but using B; using E at runtime would fail with the same error.

Solution 2: Inlining the rule on 1.11.1+

See #337, specifically, this commit.

What this does is to effectively remove the dependency on ECExt, by:

1. Inlining the output of

   @macroexpand Enzyme.@import_rrule typeof(Bijectors.find_alpha) Real Real Real
   @macroexpand Enzyme.@import_frule typeof(Bijectors.find_alpha) Real Real Real

2. Replacing occurrences of $(Expr(:meta, :inline)) in the expansion with Base.@_inline_meta.

3. Importing C from B via using Bijectors: Bijectors, ChainRulesCore.

   Note that using ChainRulesCore directly works on Julia 1.11, but fails on 1.10.6, because this requires C to be added as one of the triggers to the extension, and this bug makes it impossible for a strong dependency to be an extension trigger.

4. Adding EnzymeCore as an extension trigger. This is necessary because the expansion of the macros includes functions from EnzymeCore.

We make sure to only do this on Julia 1.11.1+; on previous versions where the loading order was not changed, we could continue to implicitly rely on the original loading order.

The result of this is that there are no errors at either precompilation or import stage.

However, the drawback is that this inlined code is specific for a particular version of Enzyme, in this case, v0.12.36. For example, in the time between that version and now, there has been one change to that extension, which is to add a config argument. To maintain compatibility with newer versions of Enzyme, we would have to follow the steps above to re-generate the code to be inlined.

Solution 3: Disabling the extension on 1.11.1+

See #337. As before, this is only done on Julia 1.11.1+, to avoid breaking compatibility with older versions.

The rationale for this is that Enzyme itself is already broken on Julia 1.11, so we don't gain anything by including the code in the extension.

Although this makes life much simpler right now, it means that we have to re-add it back when Enzyme is fixed for Julia 1.11.

Solution 4: Explicitly specifying the dependency order

This Julia PR offers a solution to this: if the extension triggers of ECExt are a strict subset of the extension triggers of BEExt, then ECExt will be loaded before BEExt.

The triggers of ECExt are ["Enzyme", "ChainRulesCore"] (the former is implicit because the extension is defined inside Enzyme). So, if we declare the triggers of BEExt to be ["Bijectors", "Enzyme", "ChainRulesCore"] (again only the last two need to be in Project.toml) then once that PR is merged we can now explicitly rely on the correct loading order.

The only drawback of this is the aforementioned bug, that we cannot have strong dependencies (like C) be an extension trigger on Julia 1.10. Thus, as it currently stands, this solution would not be backwards-compatible.

I've already dropped a comment on the Julia PR to let the maintainers know.

What did we do in Bijectors v0.14.0?

We did (3).

What's the ideal solution?

IMO, the ideal solution would be (4), BUT we would need to:

1. Ensure that the strongdep bug is fixed on Julia 1.10. This bugfix relies on this PR being backported to 1.10.
2. Drop compatibility with pre-1.10. This isn't really a problem because we've already done that on Turing.jl itself.

[wsmoses]

Enzyme now supports 1.11, so the loss of the rules for 1.11 is now problematic for users =/

[penelopeysm]

Good thing I wrote up an issue on how to deal with it! 😬

(4) is not a workable solution at present (the PR for 1.11 is merged, but 1.11.2 is not released; and the PR for 1.10 is not merged yet, and we can't version-guard it because this solution relies on changing Project.toml - we can't have one for 1.10 and one for 1.11. We could do something like releasing different versions of Bijectors for each version of Julia, but again this needs to wait for 1.11.2).

(3) is what we have now, which is a no-go. (1) is a poor solution.

Maybe it's time to move back to solution (2)?

[penelopeysm]

I think we should merge #341 first before handling this, though.

[penelopeysm]

Hmm, actually, I'm wrong; #341 and this need to be merged together. I'll open a PR to that PR.