Google Cloud just recently released "The PyTorch developer's guide to JAX fundamentals".

Contains a side-by-side implementation of a training loop in both #Pytorch and #JAX + Flax NNX sor those interested in exploring the JAX world in familiar terms.

Link: https://cloud.google.com/blog/products/ai-machine-learning/guide-to-jax-for-pytorch-developers

https://xn--r1a.website/DataScienceM