Tutorials for deep learning frameworks (PyTorch/JAX/Tensorflow)

JAX:

• Clean, NumPy-like API
• Excellent support for forward derivatives and higher-order derivatives
• Functional style, user maintains state explicitly. Avoids lots of potential bugs (especially random number generation).