TuringLang
Test against Enzyme
Currently failing, using this to run CI and catch issues. Uses code from https://github.com/TuringLang/Bijectors.jl/pull/240.
CI currently fails with
julia: /buildworker/worker/package_linux64/build/src/llvm-late-gc-lowering.cpp:872: std::vector<int, std::allocator<int> > LateLowerGCFrame::NumberAllBase(State&, llvm::Value*): Assertion `Tracked.size() == BaseNumbers.size()' failed.
Running the tests locally I also see the following
Expression: ≈(Enzyme.gradient(Enzyme.Reverse, f, x), finitediff, rtol = rtol, atol = atol)
Enzyme compilation failed due to illegal type analysis.
Current scope:
; Function Attrs: mustprogress willreturn
define internal fastcc "enzyme_type"="{[-1]:Float@double}" double @preprocess_julia__solve__70_20557({ { [3 x double] }, [4 x double], { i64, i8 } } addrspace(11)* nocapture noundef nonnull readonly align 8 dereferenceable(72) "enzyme_type"="{[-1]:Pointer, [-1,0]:Float@double, [-1,8]:Float@double, [-1,16]:Float@double, [-1,24]:Float@double, [-1,32]:Float@double, [-1,40]:Float@double, [-1,48]:Float@double, [-1,56]:Integer, [-1,57]:Integer, [-1,58]:Integer, [-1,59]:Integer, [-1,60]:Integer, [-1,61]:Integer, [-1,62]:Integer, [-1,63]:Integer, [-1,64]:Integer}" "enzymejl_parmtype"="4511001616" "enzymejl_parmtype_ref"="1" %0) unnamed_addr #99 !dbg !10218 {
[blahblahblah]
This is with Enzyme 0.12.19. @wsmoses, could you let me know if these are issues familiar to you, or if I should work to make a minimal reproduction of them?
What is the Julia version of the first error. That looks like a Julia compiler bug that was subsequently fixed
The first one is on Julia 1.6, so the compiler bug it probably is then. I'll make CI run only on latest release and retry, and try to make an MWE of the latter.
Illegal type analysis error issue here: https://github.com/EnzymeAD/Enzyme.jl/issues/1573
The Assertion failed: (vec.back() == currentBlock) issue: https://github.com/EnzymeAD/Enzyme.jl/issues/1574
@mhauru like I mentioned on the closed Enzyme issue, the illegal type analysis is thrown due to the presence of an unknown bithack in Roots.jl , rather than risk an incorrect answer: https://github.com/JuliaMath/Roots.jl/blob/4263c7683423af209af1879bd9cd223b0ba55acd/src/Bracketing/bisection.jl#L135
It looks like you all defined a custom rule for tapir/etc for find_alpha for it, which I presume is a reasonable place for a custom rule for it here
@willtebbutt @mhauru @yebai mind importing the same rule in Enzyme. Feel free to use Enzyme's chain rule import macro if you like: https://github.com/EnzymeAD/Enzyme.jl/blob/c4068fc96a9b574fd08593178611b153cb71ac5b/ext/EnzymeChainRulesCoreExt.jl#L126 If you test with both forward and reverse you'll need to import both forward and reverse rules
It looks like you all defined a custom rule for tapir/etc for find_alpha for it, which I presume is a reasonable place for a custom rule for it here
Even if Enzyme would not throw an error, one really should define a custom rule for find_alpha (it's very simple) instead of differentiating through the root finding algorithm.
More generally, similar to the ForwardDiff support IMO it would be good to add support for Enzyme to Roots.jl such that the implicit function theorem is always exploited.
Sure, but I definitely appreciate that Enzyme early errors rather than risking bad behavior here!
Since it looks like this repo already adds custom rules for this, can that be added here to start and after tests are validated possibly upstream to roots.jl
Thanks @wsmoses, I added an import of the existing ChainRule for find_alpha. This uncovered a new issue, filed here: https://github.com/EnzymeAD/Enzyme.jl/issues/1577
since it looks like all of this passes, can this get merged? @devmotion @yebai @mhauru
@wsmoses, one of the tests started failing when I enabled reverse mode Enzyme checks for Julia 1.6. Could you please check the CI logs and let me know if this is a known issue with Julia 1.6 and what the appropriate version bound would be: https://github.com/TuringLang/Bijectors.jl/actions/runs/9871431348/job/27259337160?pr=318
Hm yeah I haven't seen that one before. Open an issue and will investigate.
In interim, maybe we just enable for 1.10+ for now and update bounds once that's fixed
@wsmoses When assertation fails (like here), can we throw an error so Julia can handle it, instead of aborting/exiting? In my view, abouting a Julia session should only be used in unrecoverable scenarios. Related: https://github.com/EnzymeAD/Enzyme.jl/issues/1625
Reported the latest issue here: https://github.com/EnzymeAD/Enzyme.jl/issues/1629 It gets triggered on Julia 1.6 and 1.7, but not 1.8.
If we are happy with Enzyme only being tested on 1.10, this could be ready to merge. I'm the owner, so someone else needs to review.
or so Julia can handle it, instead of aborting/exiting? In my view, abouting a Julia session should only be used in unrecoverabl
Oh yeah totally agreed, PR's welcome!
or so Julia can handle it, instead of aborting/exiting? In my view, abouting a Julia session should only be used in unrecoverabl
Oh yeah totally agreed, PR's welcome!
@mhauru can you help create an Enzyme issue for this? Thanks!