How deeply does Deep Learning with JAX cover structured control flow?

Descends from tracing errors to fori_loop and cond implementations, Jaxpr output, static arguments, and recompilation tradeoffs.

Structured control flow

As explained in Deep Learning with JAX

Control-flow constructs in JAX that can be compiled from traces, often using lower-level lax APIs rather than ordinary Python runtime branching.

What happens when a JAX loop or branch depends on an input value? Ordinary Python control flow cannot be traced reliably. Structured control flow replaces it with jax.lax primitives that JAX can compile and differentiate. section 5.3.3 section 3.6

Use the lower-level API when:

A small set of concrete values is acceptable; mark arguments static, understanding that new values trigger recompilation. section 5.5

A traceable loop and branch

f5 = lambda x: jax.lax.fori_loop(0, x, lambda i, v: v + i, 0) relu = lambda x: jax.lax.cond(x > 0, lambda x: x, lambda x: 0.0, x) Both functions accept input-dependent values while remaining compatible with make_jaxpr() and jit().

How Deep Learning with JAX covers it

liveBook manning.com

Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

lax.fori_loopuses · A JAX control-flow primitive for looping over a fixed integer range while remaining traceable and compilable.
loopcontrasts with · A loop is a programming control-flow construct that repeatedly executes a block of instructions according to an iteration rule or continuation condition.
just-in-time compilationproduces · Just-in-time compilation translates or optimizes program code during execution, often using runtime information to generate specialized executable instructions.
conditional logiccontrasts with · Conditional logic is a programming technique that selects execution paths, actions, or result values according to evaluated conditions.
jax.laxpart of · The lower-level JAX API that provides primitive operations underlying higher-level array libraries such as jax.numpy.

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.