# \[Zygote\] Derivative of derivative compilation time
**URL:** https://discourse.julialang.org/t/zygote-derivative-of-derivative-compilation-time/95537
**Category:** Machine Learning
**Created:** [March 4, 2023, 10:42am UTC](https://discourse.julialang.org/t/zygote-derivative-of-derivative-compilation-time/95537 "2023-03-04T10:42:29Z")
**Posts on this page:** 5
**Page:** 1
### Author:  [@mariusd](https://discourse.julialang.org/u/mariusd)
#### Post date: [March 4, 2023, 10:42am UTC](https://discourse.julialang.org/t/zygote-derivative-of-derivative-compilation-time/95537/1 "2023-03-04T10:42:29Z")
In the following code I calculate derivative fx, then derivative of derivative fxy, and then fxyz.
The code works and gives the correct result, however the compilation time is huge for the fxyz derivative. Is there a way to improve this behavior?
```julia
using Zygote
# the function of interest
f(x, y, z) = x*y*z
# zygote gradients
fx(x, y, z) = gradient(x->f(x, y, z), x)[1]
fxy(x, y, z) = gradient(y->fx(x, y, z), y)[1]
fxyz(x, y, z) = gradient(z->fxy(x, y, z), z)[1]
@time fx(1.f0, 2.f0, 3.f0)
@time fxy(1.f0, 2.f0, 3.f0)
@time fxyz(1.f0, 2.f0, 3.f0)
@code_warntype fx(1.f0, 2.f0, 3.f0)
@code_warntype fxy(1.f0, 2.f0, 3.f0)
@code_warntype fxyz(1.f0, 2.f0, 3.f0)
```
Timing output:
```julia
0.000003 seconds
10.598413 seconds (21.78 M allocations: 1.099 GiB, 3.89% gc time, 99.75% compilation time)
183.757771 seconds (274.62 M allocations: 13.275 GiB, 2.83% gc time, 98.93% compilation time)
```
If we look at the warntype outputs we can see:
for fx: `Main.gradient(%6, x)::Tuple{Float32}`
for fxy: `Main.gradient(%6, y)::Union{Nothing, Tuple{Any}}`
for fxyz: `Main.gradient(%6, y)::Union{Nothing, Tuple{Any}}`
For fx the function correctly returns Tuple{Float32} but for the other we get Union{Nothing, Tuple{Any}}. Can we make fxy return Tuple{Float32} as well?
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [March 4, 2023, 3:37pm UTC](https://discourse.julialang.org/t/zygote-derivative-of-derivative-compilation-time/95537/2 "2023-03-04T15:37:27Z")
> [@mariusd](#):
>
> The code works and gives the correct result, however the compilation time is huge for the fxyz derivative. Is there a way to improve this behavior?
Maybe try [GitHub - JuliaDiff/ForwardDiff.jl: Forward Mode Automatic Differentiation for Julia](https://github.com/JuliaDiff/ForwardDiff.jl) or [GitHub - JuliaDiff/TaylorDiff.jl: Taylor-mode automatic differentiation for higher-order derivatives](https://github.com/JuliaDiff/TaylorDiff.jl) … you don’t really need reverse-mode AD for tiny numbers of parameters, and forward mode is much less demanding of the compiler.
In particular, I tried your code adapted to `ForwardDiff.derivative` and it is basically instantaneous (`@time` reports `0.000000 seconds` even on the first call).
---
### Author:  [@kristoffer.carlsson](https://discourse.julialang.org/u/kristoffer.carlsson)
#### Post date: [March 4, 2023, 4:03pm UTC](https://discourse.julialang.org/t/zygote-derivative-of-derivative-compilation-time/95537/3 "2023-03-04T16:03:52Z")
> [@stevengj](#):
>
> (`@time` reports `0.000000 seconds` even on the first call).
Note that you might need `@time @eval ...` because otherwise things might compile before the timing starts.
---
### Author:  [@mariusd](https://discourse.julialang.org/u/mariusd)
#### Post date: [March 4, 2023, 5:12pm UTC](https://discourse.julialang.org/t/zygote-derivative-of-derivative-compilation-time/95537/4 "2023-03-04T17:12:37Z")
I think I have chosen an over simplified example. I have a Flux NN and that’s why Zygote is of interest here. I will create another MWE.
Still I’m wondering if we can make the above functions compile and run faster, using only Zygote.
---
### Author:  [@mariusd](https://discourse.julialang.org/u/mariusd)
#### Post date: [March 4, 2023, 9:57pm UTC](https://discourse.julialang.org/t/zygote-derivative-of-derivative-compilation-time/95537/5 "2023-03-04T21:57:53Z")
At the end of the day, I have used the recommended ForwardDiff over Zygote approach. Here is an over-simplified MWE. This solution also works nicely for my Flux network. The Zygote over Zygote does not compile for my Flux network due to mutation errors, which can be expected according to the documentation: [Limitations · Zygote](https://fluxml.ai/Zygote.jl/stable/limitations/#Second-derivatives-1).
```julia
using ForwardDiff
using Zygote
f(x, y, z) = x*y*z
fx(x, y, z) = Zygote.gradient(x->f(x, y, z), x)[1]
a = Float32[1, 2, 3]
fx(a...)
fx_yz(x, y, z) = ForwardDiff.gradient(yz->fx(x, yz...), [y, z])
fx_yz(a...)
```