You've worked with JAX's core transformations: jit for compilation, grad for differentiation, and vmap for vectorization. While these are fundamental, many sophisticated models require more intricate computational structures, such as recurrent steps, conditional execution, or dynamic loops.

This chapter introduces JAX's functional control flow primitives, which allow you to express these complex patterns within JAX's traceable and compilable framework. We will cover:

• Using lax.scan to efficiently implement sequential operations like those found in Recurrent Neural Networks (RNNs).
• Implementing conditional logic inside compiled functions with lax.cond.
• Handling dynamic iteration using lax.while_loop.
• Understanding how these control flow operations interact with jit, grad, and vmap.
• Applying masking techniques for selective computation, often needed for variable-length data.
• How JAX stages Python code involving closures during tracing.

By the end of this chapter, you will be able to construct and reason about JAX functions that incorporate loops, conditionals, and sequential dependencies, preparing you to build more complex models and algorithms.