Up and Running with JAX - JIT Compilation, Vectorizing Transformations and autodiff

I first learned about JAX a few years back in an article on functorch. Functorch was a library that brought JAX-like composable function transformations to PyTorch, initially developed as a separate library but has since been fully integrated into PyTorch’s core (as of PyTorch 2.0). I’ve recently invested time in learning JAX, which has proven incredibly worthwhile. The clean functional approach…