Случайно попался мне на глаза, пожалуй, лучший манул по JAX, который покрывает множество тем от базовых вещей до более продвинутых вроде multi-node multi-GPU обучения.





https://www.jan-ebert.com/blog/jax-intro.html