JAX's design, centered around pure functions and transformations like jit and grad, requires specific approaches when dealing with computations that inherently involve state. Standard Python patterns often rely on modifying objects in place, which doesn't align well with JAX's functional nature, especially when code needs to be compiled or differentiated. How can we update model parameters during training, manage optimizer states, or handle recurrent computations within these constraints?

This chapter introduces functional programming patterns for managing state in JAX. You will learn about the concept of functional purity and why it matters for JAX transformations. We will cover the primary technique for handling state: explicitly passing it as input to functions and receiving updated state as output. You'll see how JAX's PyTree utility simplifies working with complex, nested state structures (like dictionaries or lists of parameters). We'll apply these ideas through practical examples, such as implementing a stateful counter and managing the state required for a simple optimization algorithm, ensuring these patterns work correctly when combined with jit, grad, and other transformations.