Zygote much slower than JAX for automatic differentiation of energy

Specific Domains Machine Learning
, , , ,
22

What about the code @gdalle posted? That should be quite fast and would work on GPUs with Zygote

23

But we understood that Zygote is very slow in this case. It works on GPU, but it is slow.

24

Then you might have to wait a few days until a solution on the Zygote side materializes. Turn on notifications for this issue:

5 Likes