Bijectors 0.13.18 not working with Enzyme on Julia 1.11.1

[miguelbiron]

Hi -- the bound for Enzyme's compat is too lose right now for Julia 1.11.1....

EDIT: Sorry for the multiple changes but now I realize that the current release of Bijectors is not working with Enzyme under Julia 1.11.1. Same code works fine on 1.10.5. Simplest reproducer:

1. Create a new Project.toml with

[deps]
Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"

[compat]
Enzyme = "<0.12.23"

Note: the compat entry is needed, otherwise the Pkg system will choose Bijectors v0.13.16 and Enzyme v0.13.11 (unsupported by Bijectors v0.13.18)
2. Run

$ julia --project -e "using Pkg; Pkg.instantiate(); Pkg.precompile()"

    Updating registry at `~/.julia/registries/General.toml`
    Updating `~/projects/temp/Project.toml`
  [76274a88] + Bijectors v0.13.18
⌅ [7da242da] + Enzyme v0.12.22
    Updating `~/projects/temp/Manifest.toml`
  [7d9f7c33] + Accessors v0.1.38
  [79e6a3ab] + Adapt v4.0.4
  [66dad0bd] + AliasTables v1.1.3
  [dce04be8] + ArgCheck v2.3.0
  [76274a88] + Bijectors v0.13.18
  [fa961155] + CEnum v0.5.0
  [082447d4] + ChainRules v1.71.0
  [d360d2e6] + ChainRulesCore v1.25.0
  [9e997f8a] + ChangesOfVariables v0.1.9
  [38540f10] + CommonSolve v0.2.4
  [34da2185] + Compat v4.16.0
  [a33af91c] + CompositionsBase v0.1.2
  [187b0558] + ConstructionBase v1.5.8
  [9a962f9c] + DataAPI v1.16.0
  [864edb3b] + DataStructures v0.18.20
  [e2d170a0] + DataValueInterfaces v1.0.0
  [31c24e10] + Distributions v0.25.112
  [ffbed154] + DocStringExtensions v0.9.3
⌅ [7da242da] + Enzyme v0.12.22
⌅ [f151be2c] + EnzymeCore v0.7.8
  [e2ba6199] + ExprTools v0.1.10
  [1a297f60] + FillArrays v1.13.0
  [d9f16b24] + Functors v0.4.12
⌅ [46192b85] + GPUArraysCore v0.1.6
⌅ [61eb1bfa] + GPUCompiler v0.26.7
  [34004b35] + HypergeometricFunctions v0.3.24
  [3587e190] + InverseFunctions v0.1.17
  [92d709cd] + IrrationalConstants v0.2.2
  [82899510] + IteratorInterfaceExtensions v1.0.0
  [692b3bcd] + JLLWrappers v1.6.1
⌅ [929cbde3] + LLVM v8.1.0
  [2ab3a3ac] + LogExpFunctions v0.3.28
  [1914dd2f] + MacroTools v0.5.13
  [dbb5928d] + MappedArrays v0.4.2
  [e1d29d7a] + Missings v1.2.0
  [d8793406] + ObjectFile v0.4.2
  [bac558e1] + OrderedCollections v1.6.3
  [90014a1f] + PDMats v0.11.31
  [21216c6a] + Preferences v1.4.3
  [43287f4e] + PtrArrays v1.2.1
  [1fd47b50] + QuadGK v2.11.1
  [c1ae055f] + RealDot v0.1.0
  [189a3867] + Reexport v1.2.2
  [ae029012] + Requires v1.3.0
  [79098fc4] + Rmath v0.8.0
  [f2b01f46] + Roots v2.2.1
  [6c6a2e73] + Scratch v1.2.1
  [a2af1166] + SortingAlgorithms v1.2.1
  [dc90abb0] + SparseInverseSubset v0.1.2
  [276daf66] + SpecialFunctions v2.4.0
  [10745b16] + Statistics v1.11.1
  [82ae8749] + StatsAPI v1.7.0
  [2913bbd2] + StatsBase v0.34.3
  [4c63d2b9] + StatsFuns v1.3.2
  [09ab397b] + StructArrays v0.6.18
  [53d494c1] + StructIO v0.3.1
  [3783bdb8] + TableTraits v1.0.1
  [bd369af6] + Tables v1.12.0
  [a759f4b9] + TimerOutputs v0.5.25
⌅ [7cc45869] + Enzyme_jll v0.0.133+0
⌅ [dad2f222] + LLVMExtra_jll v0.0.31+0
  [efe28fd5] + OpenSpecFun_jll v0.5.5+0
  [f50d1b31] + Rmath_jll v0.5.1+0
  [0dad84c5] + ArgTools v1.1.2
  [56f22d72] + Artifacts v1.11.0
  [2a0f44e3] + Base64 v1.11.0
  [ade2ca70] + Dates v1.11.0
  [8ba89e20] + Distributed v1.11.0
  [f43a241f] + Downloads v1.6.0
  [7b1f6079] + FileWatching v1.11.0
  [b77e0a4c] + InteractiveUtils v1.11.0
  [4af54fe1] + LazyArtifacts v1.11.0
  [b27032c2] + LibCURL v0.6.4
  [76f85450] + LibGit2 v1.11.0
  [8f399da3] + Libdl v1.11.0
  [37e2e46d] + LinearAlgebra v1.11.0
  [56ddb016] + Logging v1.11.0
  [d6f4376e] + Markdown v1.11.0
  [ca575930] + NetworkOptions v1.2.0
  [44cfe95a] + Pkg v1.11.0
  [de0858da] + Printf v1.11.0
  [9a3f8284] + Random v1.11.0
  [ea8e919c] + SHA v0.7.0
  [9e88b42a] + Serialization v1.11.0
  [6462fe0b] + Sockets v1.11.0
  [2f01184e] + SparseArrays v1.11.0
  [4607b0f0] + SuiteSparse
  [fa267f1f] + TOML v1.0.3
  [a4e569a6] + Tar v1.10.0
  [cf7118a7] + UUIDs v1.11.0
  [4ec0a83e] + Unicode v1.11.0
  [e66e0078] + CompilerSupportLibraries_jll v1.1.1+0
  [deac9b47] + LibCURL_jll v8.6.0+0
  [e37daf67] + LibGit2_jll v1.7.2+0
  [29816b5a] + LibSSH2_jll v1.11.0+1
  [c8ffd9c3] + MbedTLS_jll v2.28.6+0
  [14a3606d] + MozillaCACerts_jll v2023.12.12
  [4536629a] + OpenBLAS_jll v0.3.27+1
  [05823500] + OpenLibm_jll v0.8.1+2
  [bea87d4a] + SuiteSparse_jll v7.7.0+0
  [83775a58] + Zlib_jll v1.2.13+1
  [8e850b90] + libblastrampoline_jll v5.11.0+0
  [8e850ede] + nghttp2_jll v1.59.0+0
  [3f19e933] + p7zip_jll v17.4.0+2
        Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated -m`
Precompiling project...
  ✗ Bijectors → BijectorsEnzymeExt
  0 dependencies successfully precompiled in 6 seconds. 122 already precompiled.

The following 1 direct dependency failed to precompile:

BijectorsEnzymeExt 

Failed to precompile BijectorsEnzymeExt [8f0b585a-b8ef-5883-94dd-02189bbe38a1] to "/home/mbiron/.julia/compiled/v1.11/BijectorsEnzymeExt/jl_nmLqmX".
ERROR: LoadError: MethodError: no method matching _import_rrule(::Expr, ::Symbol, ::Symbol, ::Symbol)
The function `_import_rrule` exists, but no method is defined for this combination of argument types.
Stacktrace:
 [1] var"@import_rrule"(__source__::LineNumberNode, __module__::Module, args::Vararg{Any})
   @ Enzyme ~/.julia/packages/Enzyme/SiyIj/src/Enzyme.jl:1442
 [2] include
   @ ./Base.jl:557 [inlined]
 [3] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
   @ Base ./loading.jl:2790
 [4] top-level scope
   @ stdin:5
in expression starting at /home/mbiron/.julia/packages/Bijectors/ssG0T/ext/BijectorsEnzymeExt.jl:11
in expression starting at /home/mbiron/.julia/packages/Bijectors/ssG0T/ext/BijectorsEnzymeExt.jl:1
in expression starting at stdin:Precompiling project...
  ✗ Bijectors → BijectorsEnzymeExt
  0 dependencies successfully precompiled in 6 seconds. 122 already precompiled.

ERROR: The following 1 direct dependency failed to precompile:

BijectorsEnzymeExt 

Failed to precompile BijectorsEnzymeExt [8f0b585a-b8ef-5883-94dd-02189bbe38a1] to "/home/mbiron/.julia/compiled/v1.11/BijectorsEnzymeExt/jl_BjCtxT".
ERROR: LoadError: MethodError: no method matching _import_rrule(::Expr, ::Symbol, ::Symbol, ::Symbol)
The function `_import_rrule` exists, but no method is defined for this combination of argument types.
Stacktrace:
 [1] var"@import_rrule"(__source__::LineNumberNode, __module__::Module, args::Vararg{Any})
   @ Enzyme ~/.julia/packages/Enzyme/SiyIj/src/Enzyme.jl:1442
 [2] include
   @ ./Base.jl:557 [inlined]
 [3] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
   @ Base ./loading.jl:2790
 [4] top-level scope
   @ stdin:5
in expression starting at /home/mbiron/.julia/packages/Bijectors/ssG0T/ext/BijectorsEnzymeExt.jl:11
in expression starting at /home/mbiron/.julia/packages/Bijectors/ssG0T/ext/BijectorsEnzymeExt.jl:1
in expression starting at stdin:

I'm seeing this in our CI, for example https://github.com/Julia-Tempering/Pigeons.jl/actions/runs/11393614334/job/31702242507?pr=284#step:8:1112.

[penelopeysm]

Hi @miguelbiron, thanks for reporting!

I don't know exactly what causes this, but I could hazard a guess. The _import_rrule method we need isn't defined in Enzyme itself but rather in EnzymeChainRulesCoreExt: https://github.com/EnzymeAD/Enzyme.jl/blob/v0.12.36/ext/EnzymeChainRulesCoreExt.jl

Could it be that 1.10 precompiles EnzymeChainRulesCoreExt before precompiling BijectorsEnzymeExt, but somehow the order is reversed in 1.11...?

I'll try to see if there's a solution to this.

[penelopeysm]

JuliaLang/julia#56204

[miguelbiron]

Hey @penelopeysm -- yes I imagined that change would be the underlying cause of this bug. It hit us too at Pigeons.jl. And also it seems that Enzyme does not currently support 1.11, perhaps for related reasons.

[penelopeysm]

Hehe, seems that that was quite a wide-ranging change. We didn't even use get_extension anywhere and it still broke (maybe Enzyme does, I didn't look too closely).

I've just adopted the workaround suggested in that issue, so at least if we get errors with Enzyme it's because of Enzyme incompatibility with Julia 1.11, and not on our end.