I’m almost unable to compute ant gradients (EDIT 9 hrs later: meant to say “any gradients”, but got gradients for ants) in Zygote. I initially attempted to optimize my model using around a million parameters, but it got stuck computing the first gradient. I simplified the model so now it had only 40…

Enzyme v0.13.28. Works on 1.10 not on 1.11

I’ve also encountered a few more bugs with Enzyme on 1.11 than on 1.10. And its behavior can be platform-dependent since the LLVM codegen is different

We’re on the nighmare of wsmoses

Yeah the transition to 1.11 has been a big undertaking for Enzyme due to the introduction of Memory and rebasing Array on top of it. You should definitely submit a github issue with this example.

But for the love of god don’t submit any example involving DI. Craft a pure-Enzyme MWE

Confirming that you get a runtime activity error on 1.11, and the workarounds are to either make weights_ctx Duplicated as above, or change the mode to Enzyme.set_runtime_activity(Enzyme.Reverse). The latter seems to be a hair faster.

I made the mwe using Enzyme function foo(x,y) y2 = reshape(y,2,5) return sum(x .+ y2[:]) end x = rand(10) y = rand(10) dx = Enzyme.make_zero(x) Enzyme.autodiff(Enzyme.Reverse, foo, Enzyme.Duplicated(x,dx), Enzyme.Const(y)) dx I file the issue, it seems like it’s not only about reshape, y…

Looks like it’s already reported here, at least for reshape: error when differentiating `reshape` · Issue #2214 · EnzymeAD/Enzyme.jl · GitHub The y.^2 issue hits a different codepath (broadcasting machinery, not GenericMemory), so I suppose it might warrant a separate issue.

ok it’s the one i used anyway

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

Specific Domains Optimization (Mathematical)

yolhan_mannes January 30, 2025, 8:57pm 43

Yes I did that twice I think, I will avoid it promise

Topic		Replies	Views
Zygote much slower than JAX for automatic differentiation of energy Machine Learning performance , flux , zygote , jax , lux	22	1749	May 15, 2024
Any faster way of computing small gradients? Performance zygote , forwarddiff , symbolics , autodiff	21	2143	August 11, 2022
Speeding up Zygote autodiff for numerical loop Performance question	13	404	December 16, 2024
Lux, ComponentArrays and flat parameters : computing the gradient works with Zygote but not with Enzyme New to Julia enzyme	16	1895	May 14, 2024
Performance comparison - Flux.jl's Adam vs Jax's Adam Performance question , package , performance , flux	33	4147	November 3, 2022

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

Related topics