That's incorrect.
If you work with mid-sized neural networks and MCMC sampling, allocations start to play a significant role (And Flux.jl is bad at preallocation).
Prealloc.jl does not work properly. Zygote.jl adds even more allocations to the mix...
Jax/XLA completely solves this problem. Yes, it's annoying that you have to work with a static graph but if your problem fits the description... it's great.
Jax/XLA completely solves this problem. Yes, it's annoying that you have to work with a static graph but if your problem fits the description... it's great.