Test against Enzyme

[mhauru]

Currently failing, using this to run CI and catch issues. Uses code from https://github.com/TuringLang/Bijectors.jl/pull/240.

[mhauru]

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?

[wsmoses]

What is the Julia version of the first error. That looks like a Julia compiler bug that was subsequently fixed

[wsmoses]

The latter issue is not known please open an issue with a mwe

[mhauru]

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.

[mhauru]

Illegal type analysis error issue here: https://github.com/EnzymeAD/Enzyme.jl/issues/1573

[mhauru]

The Assertion failed: (vec.back() == currentBlock) issue: https://github.com/EnzymeAD/Enzyme.jl/issues/1574

[wsmoses]

@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

[devmotion]

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.

[wsmoses]

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

[mhauru]

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

[wsmoses]

since it looks like all of this passes, can this get merged? @devmotion @yebai @mhauru