Introduction to JAX

#Learning Objectives

• Write pure-function neural network code using JAX's functional API (jax.numpy, jax.grad, jax.jit, jax.vmap)
• Explain the key design difference between PyTorch's eager mutation and JAX's functional compilation model
• Apply jit compilation and vmap vectorization to accelerate training loops compared to naive Python
• Train a simple network in JAX and contrast the explicit state management with PyTorch's object-oriented approach