On first looking into JAX

This blog post provides an in-depth comparison between JAX and PyTorch through the author's personal experience in porting a neural network model from PyTorch to JAX. The author shares insights on the fundamental differences in philosophy between the two frameworks, discussing aspects like functional programming, JIT compilation, and the pros and cons of each. The latter part of the post explores potential downsides of JAX and raises questions about its long-term viability as a tool. Ultimately, the author expresses a mix of fascination and skepticism regarding JAX and its approach to software design.