I’ve got a Turing Bayesian model with hundreds of thousands of data points (as usual 
). For whatever reason with my current situation is that LBFGS is working to optimize it. I tried to speed up with AutoReverseDiff(true) and it crashed Julia with out of memory. I thought maybe I’d try AutoEnzyme() so I’m installing it, but I’m wondering if Enzyme is ready for cautious everyday use at the moment or is it still in a very experimental not at all reliable kind of stage of development?
1 Like
Enzyme.jl on Julia 1.10 passes all of the Turing.jl test suite tests (though there are some weird distributions.jl things that may come up).
It definitely doesn’t mean that every code ever will work out of the box the first time, but at least for your use case I think it should be reasonable exepectations of success (and if something goes awry open an issue).
I will note that you should definitely use Julia 1.10.5 since there were bugfixes to Julia’s garbage collection that were causing segfaults, fixed in that patch.
7 Likes
This is good to know. I did try it this morning, and I got a cryptic message about was I using a constant global for storage or some such thing… I’ll post the error in a bit.
Right now I’m on 1.10.4 so I’ll upgrade.
Oh yeah for sure, if you can file an issue with whatever you hit, we can try to get it fixed quickly (and also try to improve to create a more clear error message)!
I’ve been using Enzyme quite a bit lately. It’s been a game-changer for the kind of mutation-heavy code I write. In my experience, it usually works for reasonably sane and idiomatic Julia code, however, if your calculations are large and complex, chances are you’ll occasionally hit an edge case that’s not yet supported. However, the amount of supported Julia code is rapidly expanding, and you might find that the modifications required to make your code work with Enzyme are improvements anyway (though it can sometimes be tricky to figure out exactly which change you need to make).
Also, when I’ve reported issues, @wsmoses has been incredibly responsive, and more often than not it’s fixed within a day or two. (This is not to establish a burdensome expectation, priorities may change, and open-source developers don’t owe any of us anything—just voicing my appreciation.)
In short, I recommend giving it a try, and while I’m at it, let me give a big shoutout to @wsmoses, @vchuravy, and the other contributors—it’s a fantastic, yet rapidly improving piece of software.
The workaround to try in this case, until the issue is fixed within Enzyme, is to add Enzyme.API.runtimeActivity!(true) right after importing Enzyme. But also, take a close look and see if some variable marked Const is mutated by the function you’re differentiating (if you’re using AutoEnzyme() in a high-level interface, this might require opening the black box and understanding a little bit about the Enzyme API).
11 Likes
I am in fact using AutoEnzyme() in Turing.maximum_a_posteriori with the LBFGS() solver. I will try the workaround.
You probably need to do Enzyme.API.runtimeActivity!(true) at the top of the file after using Enzyme
that is, assumign your error message is:
ERROR: Mismatched activity for: store {} addrspace(10)* %.fca.0.extract, {} addrspace(10)* addrspace(10)* %.repack, align 8, !dbg !44, !tbaa !47, !alias.scope !51, !noalias !52 const val: %.fca.0.extract = extractvalue [2 x {} addrspace(10)*] %0, 0, !dbg !8
value=Unknown object of type Matrix{Float64}
llvalue= %.fca.0.extract = extractvalue [2 x {} addrspace(10)*] %0, 0, !dbg !8
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now
[which says in the error message to “use Enzyme.API.runtimeActivity!(true) as a workaround”]
The error did in fact say that, and I used it and still got some problems, but I didn’t invoke it immediately after loading Enzyme. So now I’ve just tried doing that, and it seems to work! At least, the minimization is happening without an immediate error. Haven’t seen the results yet.
Spoke too soon… here’s the output:
ERROR: Enzyme execution failed.
Enzyme: Augmented forward pass custom rule Tuple{EnzymeCore.EnzymeRules.ConfigWidth{1, true, false, (false, true)}, Const{typeof(deepcopy)}, Type{Const{Vector{Union{Missing, Int64}}}}, Const{Vector{Union{Missing, Int64}}}} return type mismatch, expected EnzymeCore.EnzymeRules.AugmentedReturn{Vector{Union{Missing, Int64}}, Nothing, Any} found EnzymeCore.EnzymeRules.AugmentedReturn{Vector{Union{Missing, Int64}}, Vector{Union{Missing, Int64}}, Vector{Union{Missing, Int64}}}
Stacktrace:
[1] macro expansion
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:7099 [inlined]
[2] enzyme_call
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6708 [inlined]
[3] AugmentedForwardThunk
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6596 [inlined]
[4] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(deepcopy), df::Nothing, primal_1::Vector{…}, shadow_1_1::Nothing)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/YWQiS/src/rules/jitrules.jl:338
[5] matchingvalue
@ ~/.julia/packages/DynamicPPL/DvdZw/src/compiler.jl:744 [inlined]
[6] matchingvalue
@ ~/.julia/packages/DynamicPPL/DvdZw/src/compiler.jl:764 [inlined]
[7] matchingvalue
@ ~/.julia/packages/DynamicPPL/DvdZw/src/compiler.jl:761 [inlined]
[8] macro expansion
@ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:1003 [inlined]
[9] make_evaluate_args_and_kwargs
@ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:981
[10] _evaluate!!
@ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:972 [inlined]
[11] evaluate_threadsafe!!
@ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:962 [inlined]
[12] evaluate!!
@ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:892 [inlined]
[13] LogDensityFunction
@ ~/.julia/packages/Turing/axR9Q/src/optimisation/Optimisation.jl:145
Just my two cents here. I’ve used Enzyme recently in some private projects and it worked well.
1 Like
With any AD system, I assume that a sufficiently complex program will run into cases that the AD system cannot handle (or handles inefficiently because it misses some structure). So I always want a well-documented “escape hatch” to write efficient custom rules.
On writing custom Enzyme rules, it seems to be mostly there. ChainRules.jl is a reasonably well-documented way to write custom rules, and Enzyme allows you to import a ChainRules rrule, but it warns you that the performance will be suboptimal. And it has its own API for writing efficient custom rules, but at this point it’s somewhat underdocumented. (For example: @wsmoses was kind enough to help us add custom AD rules to QuadGK recently, but the implementation used quite a few functions that aren’t well documented; it would have been quite difficult to write without his collaboration. To be fair, however, our case was more complicated than a typical user’s because it involves a higher-order function.)
3 Likes
Can you open an issue with a MWE of that?
Per the error message, it looks like the custom rule for deepcopy didn’t return the right type in your case. It should be a pretty quick fix (that also anyone else who wants to use would benefit from too).
As a relative marker, Enzyme can now autodiff through OrdinaryDiffEq explicit methods (i.e. non-adjoint direct adjoint). Of course that flexes a lot of Julia, DataStructures, ranges, logging, rules around bit hacks, etc. so it’s a pretty good milestone. There’s still some issues with the implicit algorithms because of nesting Enzyme with PreallocationTools (due to tag reinterpret) that will get fixed by using Enzyme as the Jacobian tooling. We’ve also been doing this change to ODINN.jl, working through making Enzyme handle PyCall xarrays and everything that shows up in a real model code. Those two codes are pretty hard testbeds and I think we’re on track to get them working by the end of the year.
7 Likes
The only issue is making the MWE. If I can find a minimal way to do it yeah, if not it’s downloading large quantities of Census data and soforth. I’ll make some effort though.
@ChrisRackauckas, that is FANTASTIC news. Thanks for that info.
Well, higher-order functions aren’t that unusual in Julia. Improved documentation/API for writing such rules is one of the things I’ve been wishing for too. But my use case was in fact a QuadGK call, so once again, the Enzyme folks have solved my problem.
Ah if that’s the case, I just made this which might fix it? Handle deepcopy of constant by wsmoses · Pull Request #1765 · EnzymeAD/Enzyme.jl · GitHub
Depending on what the error is, we might need a more minimal example to figure out what’s happening and why.
In this case the error message contained most of the info we need (though I’m still confused why this case would ever arise).
1 Like
Yeah we definitely need more docs on this. Always open to PR’s 
That said if you’re writing a rule and having issues, feel free to ping me. And then potentially you can help turn whatever you learn into docs?
4 Likes