Core Program Patterns

Use these recipes for traced Python programs, pytrees, batching, control flow, and dependency-aware execution.

• Vectorize Inputs with in_axes
  • Broadcast One Input
  • Batch a Nested Input
  • Pair Two Batches
• Use Object-Oriented Modules
  • Define the Module
  • Trace the Module Call
  • Batch Module Configurations
  • Receive Module-Shaped Feedback
  • Keep Runtime Boundaries Separate
• Run an LM Pipeline Concurrently
• Use Control Flow Inside a Traced Function
  • Route with switch
  • Repeat with while_loop
  • Block Feedback with stop_gradient
  • Force Ordering with depends