# Zygote much slower than JAX for automatic differentiation of energy

**URL:** https://discourse.julialang.org/t/zygote-much-slower-than-jax-for-automatic-differentiation-of-energy/114239
**Category:** Machine Learning
**Tags:** performance, flux, zygote, jax, lux
**Created:** [May 14, 2024, 11:48am UTC](https://discourse.julialang.org/t/zygote-much-slower-than-jax-for-automatic-differentiation-of-energy/114239 "2024-05-14T11:48:32Z")
**Posts on this page:** 1
**Showing post:** 3

<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: [May 14, 2024, 12:39pm UTC](https://discourse.julialang.org/t/zygote-much-slower-than-jax-for-automatic-differentiation-of-energy/114239/3 "2024-05-14T12:39:57Z")

</div>

~~If we’re looking at the allocations themselves, I don’t observe the same as you.  
On my laptop, it’s 4 Mb rather than 4 Gb.~~ Can you share the versions of Julia and the packages that you are using?

As a side note, allocations of 4 Gb would be coherent with a dense version of the sparse `H` matrix:

```julia
julia> @benchmark ones(Float32, 2^15, 2^15)
BenchmarkTools.Trial: 4 samples with 1 evaluation.
 Range (min … max): 1.505 s … 1.775 s ┊ GC (min … max): 0.02% … 14.93%
 Time (median): 1.647 s ┊ GC (median): 6.40%
 Time (mean ± σ): 1.643 s ± 133.053 ms ┊ GC (mean ± σ): 7.24% ± 7.88%

  █ █ █ █  
  █▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁█ ▁
  1.5 s Histogram: frequency by time 1.78 s <

 Memory estimate: 4.00 GiB, allocs estimate: 2.

```

I wonder if this is related to a naive chain rule which allocates the full matrix during a pullback

> [@How do you speed up the linear sparse solver in Zygote?](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801/):
>
> I have confirmed that the standard Julia solver (backslash) can be used in the automatic differentiation of Zygote. Below is a sample program. (This program calculates the gradient.) However, while the forward computation of the evaluation function is fast, the computation of the automatic derivative is very slow. (Memory is not that large.) What are you trying to do to speed up Zygote’s automatic differentiation? Is it common to write your own rrule and have it calculate the derivative dire…

Maybe @mohamed82008 will have an opinion.

---

_[View the full topic](https://discourse.julialang.org/t/zygote-much-slower-than-jax-for-automatic-differentiation-of-energy/114239)._