Control Flow

Structured control-flow transformations that are JIT-compilable and state-aware, covering conditional branching, bounded and dynamic loops, and result-collecting iteration.

Conditions

Control flow transformations that enable conditional execution of different computation branches based on runtime conditions. These functions provide efficient, JIT-compilable alternatives to Python’s native if/elif/else statements, ensuring optimal performance in compiled code.

cond(pred, true_fun, false_fun, *operands)	Conditionally apply true_fun or false_fun.
switch(index, branches, *operands)	Apply exactly one branch from branches based on index.
ifelse(conditions, branches, *operands[, ...])	Represent multi-way if/elif/else control flow.

For Loop

Transformations for structured iteration with result collection. These functions provide efficient ways to perform repeated computations while accumulating results into arrays, with optional checkpointing for memory-efficient training of deep networks.

scan(f, init, xs[, length, reverse, unroll, ...])	Scan a function over leading array axes while carrying along state.
checkpointed_scan(f, init, xs[, length, ...])	Scan a function over leading array axes while carrying along state.
for_loop(f, *xs[, length, reverse, unroll, pbar])	for-loop control flow with State.
checkpointed_for_loop(f, *xs[, length, ...])	for-loop control flow with State with a checkpointed version, similar to for_loop().
ProgressBar([freq, count, desc])	A progress bar for tracking the progress of a jitted for-loop computation.

While Loop

Dynamic iteration transformations that continue execution based on runtime conditions. These functions enable loops with variable iteration counts, essential for adaptive algorithms and convergence-based computations.

while_loop(cond_fun, body_fun, init_val)	Call body_fun repeatedly in a loop while cond_fun is True.
bounded_while_loop(cond_fun, body_fun, ...)	While loop with a bound on the maximum number of steps.