# Picking an AD Backend and Enzyme Errors

**URL:** <https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075>\
**Category:** Optimization (Mathematical)\
**Tags:** optimization, enzyme\
**Created:** [March 2, 2024, 11:44pm UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075 "2024-03-02T23:44:37Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![kiranshila](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/kiranshila/32/25475_2.png) [@kiranshila](https://discourse.julialang.org/u/kiranshila)\
**Post date:** [March 2, 2024, 11:44pm UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/1 "2024-03-02T23:44:37Z")

</div>

I’m working on an optimization problem for a paper of the form  
`cost(θ) => scalar`, where θ is potentially large, maybe 100-200 dimensional.

I was under the assumption reverse mode AD is appropriate for this task, however benchmarking a few:

For θ in R^100

For forward mode:

- `ForwardDiff` takes ~200ms
- `PolyesterForwardDiff` takes ~20ms (with 10 Chunks)

For reverse mode:

- `ReverseDiff` takes ~2000ms with significantly more allocations
- `Zygote` fails because I do quite a lot of mutation to avoid allocations
- `Diffractor` gives some type error that I don’t understand
- `Enzyme` gives [this](https://gist.githubusercontent.com/kiranshila/4cc8ccf08b3328421df25ca65564e218/raw/0edc83a240abb34c11bac9f6938ca857fbf415d0/error), which…I’m really not sure what to do with

For completeness, I opened an Enzyme issue [here](https://github.com/EnzymeAD/Enzyme.jl/issues/1319).

So, what is the best approach here? `ForwardDiff` is working, `PolyesterForwardDiff` gives a nice 10x speedup, but seem unsupported as an AD backend for SciML’s `Optimization`. And only `ReverseDiff` works in reverse-land.

I’d appreciate any insight!

---

<div class="post-metadata">

**Author:** ![kiranshila](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/kiranshila/32/25475_2.png) [@kiranshila](https://discourse.julialang.org/u/kiranshila)\
**Post date:** [March 2, 2024, 11:46pm UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/2 "2024-03-02T23:46:21Z")

</div>

For more context, the problem setup is [here](https://github.com/kiranshila/WBLNAOptim/blob/main/scripts/multi_point_optim.jl)

---

<div class="post-metadata">

**Author:** ![kiranshila](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/kiranshila/32/25475_2.png) [@kiranshila](https://discourse.julialang.org/u/kiranshila)\
**Post date:** [March 3, 2024, 12:11am UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/3 "2024-03-03T00:11:26Z")

</div>

Ah, also my mistake, `AutoPolyesterForwardDiff` is indeed apart of `Optimization`. It just hasn’t made it to the docs yet.

EDIT: However, I could not get it to work.

---

<div class="post-metadata">

**Author:** ![odow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/odow/32/28685_2.png) [@odow](https://discourse.julialang.org/u/odow)\
**Post date:** [March 3, 2024, 3:02am UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/4 "2024-03-03T03:02:23Z")

</div>

Can you post a reproducible example of your code?

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [March 3, 2024, 3:11am UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/5 "2024-03-03T03:11:08Z")

</div>

Replied on GitHub re Enzyme error.

Looks like Julia’s complex sqrt has a bit hack. We should just define a rule for it. That should fix it, but I don’t myself have time for roughly a week before I can work on it.

If you or others are interested in adding it lmk.

Doing so properly requires writing the complex sqrt here ([Enzyme/enzyme/Enzyme/InstructionDerivatives.td at 0b621884bc531329095d202f042f6599a86614ec · EnzymeAD/Enzyme · GitHub](https://github.com/EnzymeAD/Enzyme/blob/0b621884bc531329095d202f042f6599a86614ec/enzyme/Enzyme/InstructionDerivatives.td#L834) this is the complex 1/z rule) and for Julia here ([Enzyme.jl/src/compiler/interpreter.jl at 8784d1f79cf9e84028bc04c7455493d1b9dcbd31 · EnzymeAD/Enzyme.jl · GitHub](https://github.com/EnzymeAD/Enzyme.jl/blob/8784d1f79cf9e84028bc04c7455493d1b9dcbd31/src/compiler/interpreter.jl#L97)) and here ([Enzyme.jl/src/compiler.jl at 8784d1f79cf9e84028bc04c7455493d1b9dcbd31 · EnzymeAD/Enzyme.jl · GitHub](https://github.com/EnzymeAD/Enzyme.jl/blob/8784d1f79cf9e84028bc04c7455493d1b9dcbd31/src/compiler.jl#L71))

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [March 3, 2024, 4:16am UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/6 "2024-03-03T04:16:53Z")

</div>

For performant reversediff you would want to compile the tape [API · ReverseDiff.jl](https://juliadiff.org/ReverseDiff.jl/dev/api/#ReverseDiff.compile)

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [March 3, 2024, 8:09am UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/7 "2024-03-03T08:09:48Z")

</div>

> [@kiranshila](#):
>
> `Diffractor` gives some type error that I don’t understand

For completeness, at the moment Diffractor.jl is only a forward mode backend, and still rather experimental

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [March 3, 2024, 11:38am UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/8 "2024-03-03T11:38:19Z")

</div>

> [@kiranshila](#):
>
> I was under the assumption reverse mode AD is appropriate for this task, however benchmarking a few:
> 
> For θ in R^100

Is there a reason you think reverse mode is more appropriate for this task? We have done many detailed measurements in different domains, and for example [A Comparison of Automatic Differentiation and Continuous Sensitivity Analysis for Derivatives of Differential Equation Solutions | IEEE Conference Publication | IEEE Xplore](https://ieeexplore.ieee.org/abstract/document/9622796) this mentions that due to various limitations of reverse mode that 100 is the cutoff to where you should start considering reverse-mode AD.

- So anything below 100, definitely forward.
- Way above 100, definitely reverse
- Around 100? Murky waters.

And that’s without considering PolyesterForwardDiff, which using parallelism on a high core computer definitely pushes that range up a bit. Though newer versions of Enzyme definitely push that range back down a bit, so 100 is still about where I would expect it to be.

All I mean to say from this is, our benchmarks continually tell us that no, I don’t know why you would expect miracles from reverse mode here. PolyesterForwardDiff is likely to be the fastest, I would not expect even compiled tape ReverseDiff to be faster, Enzyme is the only thing that has a chance but that’s dependent on many factors (and it’s battling uphill without a multithreaded form), and you’re likely to not see a real gain after doing a bunch of work to get reverse mode to work better. I would set the expectations there and consider this an academic exercise to reverify it. If you were talking about 1000 I’d say differently, but 100 is about where the cutoff is and the multithreaded versions really make that level not “big”.

---

<div class="post-metadata">

**Author:** ![kiranshila](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/kiranshila/32/25475_2.png) [@kiranshila](https://discourse.julialang.org/u/kiranshila)\
**Post date:** [March 3, 2024, 10:24pm UTC](https://discourse.julialang.org/t/picking-an-ad-backend-and-enzyme-errors/111075/9 "2024-03-03T22:24:29Z")

</div>

Thanks for your input, Chris! What you say makes total sense.

> Is there a reason you think reverse mode is more appropriate for this task?

Well, that’s why I benchmarked it 🙂.

> And that’s without considering PolyesterForwardDiff, which using parallelism on a high core computer definitely pushes that range up a bit. Though newer versions of Enzyme definitely push that range back down a bit, so 100 is still about where I would expect it to be.

I’m excited to see how Enzyme will perfom once we fix that bug from above w.r.t complex sqrt, but otherwise I’m quite satisfied with the performance of PolyesterForwardDiff, however I haven’t quite gotten it to work in Optimization. Nothing in the docs mention PolyesterForwardDiff ([OptimizationFunction · Optimization.jl](https://docs.sciml.ai/Optimization/stable/API/optimization_function/#Automatic-Differentiation-Construction-Choice-Recommendations)), although `AutoPolyesterForwardDiff` seems to exist.

I’m not sure if this warrants opening an issue but as an example:

```julia
using Optimization, OptimizationOptimJL, PolyesterForwardDiff
rosenbrock(u, p) = (p[1] - u[1])^2 + p[2] * (u[2] - u[1]^2)^2
u0 = zeros(2)
p = [1.0, 100.0]
optf = OptimizationFunction(rosenbrock, Optimization.AutoPolyesterForwardDiff())
prob = OptimizationProblem(optf, u0, p, lb = [-1.0, -1.0], ub = [1.0, 1.0])

```

gives

```julia
ArgumentError: The passed automatic differentiation backend choice is not available. Please load the corresponding AD package PolyesterForwardDiff.

```

even though it is clearly loaded.