That's right, but at the same time Flux, the default choice is not mature at all compared to JAX. So that advantage is only on paper. I would really prefer to

That's right, but at the same time Flux, the default choice is not mature at all compared to JAX. So that advantage is only on paper. I would really prefer to use Julia, I actually dislike Python, but it's unrealistic to implement large architectures in Flux. It is buggy, lots of gradient calculations are unreliable. See some discussion here for instance: https://kidger.site/thoughts/jax-vs-julia https://kidger.site/thoughts/jax-vs-julia