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