# \`Zygote.gradient\` is 54000 TIMES slower than \`jax.gradient\`

**URL:** <https://discourse.julialang.org/t/zygote-gradient-is-54000-times-slower-than-jax-gradient/125396>\
**Category:** Optimization (Mathematical)\
**Tags:** zygote, jax\
**Created:** [January 30, 2025, 4:23pm UTC](https://discourse.julialang.org/t/zygote-gradient-is-54000-times-slower-than-jax-gradient/125396 "2025-01-30T16:23:22Z")\
**Posts on this page:** 1\
**Showing post:** 8

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [January 30, 2025, 6:31pm UTC](https://discourse.julialang.org/t/zygote-gradient-is-54000-times-slower-than-jax-gradient/125396/8 "2025-01-30T18:31:33Z")

</div>

I haven’t checked this carefully, but I think that we-writing this not to make slices at all can get us to 2ms:

```julia
julia> function loss(
               Wmid, Wctx,
               tok_mid, tok_ctx, x
       )
       tmp = sum(Wmid[tok_mid, :] .* Wctx[tok_ctx, :]; dims=2) |> vec
       -mean(@. x * log(logistic_sigmoid(tmp)) + (1 - x) * log(1 - logistic_sigmoid(tmp)))
       end
loss (generic function with 1 method)

julia> grad = @btime train(Random.Xoshiro(5), 100277, 100, 2);
  1.892 ms (274 allocations: 4.62 MiB)

julia> using Tullio # first way I thought of

julia> function loss(
               Wmid, Wctx,
               tok_mid, tok_ctx, x
       )
       @tullio tmp[k] := Wmid[tok_mid[k], c] * Wctx[tok_ctx[k], c] # sum over c
       -mean(@. x * log(logistic_sigmoid(tmp)) + (1 - x) * log(1 - logistic_sigmoid(tmp))) # sum over k
       end
loss (generic function with 1 method)

julia> grad = @btime train(Random.Xoshiro(5), 100277, 100, 2);
  1.916 ms (348 allocations: 4.62 MiB)

```

All definitions of `loss` give me zero gradient, so there could be mistakes.

---

_[View the full topic](https://discourse.julialang.org/t/zygote-gradient-is-54000-times-slower-than-jax-gradient/125396)._