Core Foundation · Local CPULaptop CPULesson 4 of 4 in Phase 03

Phase 03 / State, randomness & control flow

Branches with lax.cond

Revision v0.7.0 · 2026-10-20

Lesson · about 60 minutes · CPU practice

Implement a data-dependent branch that works inside a compiled function.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Compiled loops with lax.scan
Hardware: CPU

Download notebook

Other downloads

Python worked examples Lesson text

Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder

py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment

python phases/03-state/04-branches-with-lax-cond/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/03-state/04-branches-with-lax-cond/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.

For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:

macOS / Linux terminal · activated course environment

python -m pip install jupyterlab
python -m jupyter lab

PowerShell command · repository folder

.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter lab

Select the downloaded state-04.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.

Check: $3.0$

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 03.04 Stage 1 of 5 · Understand

1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence

What you will be able to do

Locate the decision boundary
Trace both branches and check their result types
Choose cond or where by the shape of the question
Inspect boundary behavior and derivatives

The problem

A Python branch can choose a path when it knows a value. Inside jit, that value may not be available yet. We’ll use lax.cond to make the choice at runtime, check what both branches must return, and compare a single decision with elementwise selection.

The idea

A data-dependent branch needs a control-flow operation that can represent both alternatives in a transformed program. With a scalar predicate, lax.cond selects the result of one branch while requiring compatible result structures from both.

Trace both possibilities, execute the selected rule

Consider an absolute-value rule: return the input when it is nonnegative and its negation otherwise. Both branches return the same shape and dtype, even though only one scalar case determines a particular result. The function is continuous at zero, but its mathematical derivative is not uniquely defined there.

The plot of magnitudes shows what the branch computes, not the tracing process. During tracing, Python code in both branch functions can be visited. Avoid using Python side effects as evidence that one numerical branch executed on the device.

After batching or other transformations, the implementation of control flow can change; for example, batching a conditional can turn it into selection. Keep claims about execution scope tied to the actual transformed program.

\mathcal{T} = (\text{Leaves}(\mathcal{T}),\; \text{TreeDef}(\mathcal{T})), \qquad \operatorname{map}(g, \mathcal{T})_k = g(\mathcal{T}_k)

Locate the decision boundary

The predicate $x\ge0$ is numerical data when $x$ is a traced argument. A Python if must choose a branch while tracing, before that runtime value is available. lax.cond represents the choice inside the staged computation instead.

Its arguments are a scalar Boolean predicate, two pure branch functions and an operand. Each branch receives that operand; each must return the same kind of result. Read the contract as an ordinary state transition with two possible implementations, not as a magic exception to tracing.

predicate: x>=0
true:  operand x → x
false: operand x → −x
joined result: scalar float

Trace both branches and check their result types

For a scalar predicate, cond selects one branch at runtime, but both branches are traced when building the computation. Therefore an invalid operation in a branch can fail during tracing even if a particular test input would not choose it.

Both result pytrees must agree and corresponding array leaves must have compatible shapes and dtypes. A scalar in one branch and a vector in the other cannot be joined by this contract. This is different from numerical correctness: matching output types can still hide a wrong branch formula.

Choose cond or where by the shape of the question

cond makes one scalar decision for a computation. where chooses values elementwise from compatible arrays. A vector predicate naturally belongs to where or a deliberately batched conditional.

When cond is transformed with vmap over predicates, it can become selection. Do not use a branch as a promise that side effects or invalid arithmetic will never be evaluated under transformations. Keep functions pure and design numerical expressions that are valid in the relevant domain.

Inspect boundary behavior and derivatives

Our magnitude function equals absolute value. Away from zero its derivative is + $1$ for positive inputs and $-1$ for negative inputs. At zero the mathematical function is not differentiable; a selected branch rule is a convention rather than proof of a derivative.

Test negative, positive and exact-boundary inputs because a > versus >= change can alter boundary behavior in other functions. Keep a separate argument about mathematical smoothness when differentiating conditional computations.

Both branches form one magnitude function

Predict before looking: What happens on each side of zero?

Read the figure

The horizontal axis varies the input and the vertical axis gives the returned magnitude. The V-shaped line passes through $(-3,3)$ , $(0,0)$ , and $(3,3)$ . Equal distances to the left and right of zero produce equal outputs.

The left arm comes from negating a negative input. The right arm comes from returning a nonnegative input unchanged. At zero the two values meet, so there is no jump in the output.

Connect it to the computation

The branches together implement $|x|$ . Moving right along the negative side decreases the result with slope $-1$ ; moving right along the positive side increases it with slope $1$ . The sharp corner records that change of rule.

Continuity at the join does not imply differentiability: the classical derivative is undefined at zero because the two slopes disagree. A branch-selected autodiff value there must not be confused with a derivative of a smooth curve. The picture helps separate correct branch outputs from assumptions about gradients.

Open full-size figure Download data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Where to put & run this code: Work in phases/03-state/04-branches-with-lax-cond/code/main.py (or downloaded exercises/state-04.py / exercises/state-04.ipynb) and execute from the repository root with python3 phases/03-state/04-branches-with-lax-cond/code/main.py.

1 Prepare the inputs

Create main.py in your lesson workspace. Add this first block; use the environment from setup.

main.py · append build step 1

# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

These explicit inputs define the case that the later checks will verify.

2 Build the computation

Append this block below the inputs in the same file.

main.py · append build step 2

# Step 2 — Build the computation: Both branches return the same scalar type.
# Define and JIT-compile `magnitude(x)` so XLA traces and fuses the operations:
@jax.jit
def magnitude(x):
    return jax.lax.cond(x >= 0., lambda v: v, lambda v: -v, x)

Code building blocks in this step

jax.jit(fn) / @jax.jit — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

Both branches return the same scalar type. cond takes the runtime predicate and passes $x$ to the selected numerical branch.

3 Run and check the result

Append the checks, save main.py, and run python main.py from this folder using your course environment.

main.py · append build step 3

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print(float(magnitude(-3.)))
assert jnp.allclose(magnitude(-3.), 3.)
assert jnp.allclose(magnitude(2.), 2.)

Code building blocks in this step

jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

Compare the output to the expected result below before making the exercise change.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

# Step 2 — Build the computation: Both branches return the same scalar type.
# Define and JIT-compile `magnitude(x)` so XLA traces and fuses the operations:
@jax.jit
def magnitude(x):
    return jax.lax.cond(x >= 0., lambda v: v, lambda v: -v, x)

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print(float(magnitude(-3.)))
assert jnp.allclose(magnitude(-3.), 3.)
assert jnp.allclose(magnitude(2.), 2.)

Run the example

python3 phases/03-state/04-branches-with-lax-cond/code/main.py

About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example

main.py · complete Python example

# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp
# Step 2 — Build the computation: Both branches return the same scalar type.
# Define and JIT-compile `magnitude(x)` so XLA traces and fuses the operations:
@jax.jit
def magnitude(x):
    return jax.lax.cond(x >= 0., lambda v: v, lambda v: -v, x)
# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print(float(magnitude(-3.)))
assert jnp.allclose(magnitude(-3.), 3.)
assert jnp.allclose(magnitude(2.), 2.)

Expected result: $3.0$

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:45:31.964561+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

3.0
3.0
Expected traced Python Boolean failure
Expected branch shape mismatch
PASS: state-04

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Reproduce the Python truth-test failure

Reproduce the Python truth-test failure · Python experiment

# Experiment — Reproduce the Python truth-test failure: The repair changes where the decision is represented.
def eager_branch(v):
    if v>=0.:
        return v
    return -v
assert eager_branch(-3.)==3.
# Run the boundary check and catch the expected exception:
try:
    jax.jit(eager_branch)(jnp.array(-3.))
except jax.errors.TracerBoolConversionError:
    print("Expected traced Python Boolean failure")
else:
    raise AssertionError("Expected branch error")
for v in [-3.,0.,2.]:
    assert jnp.allclose(magnitude(v),abs(v))

Expected: The traced Python branch fails; cond gives the correct values for all three cases.

The repair changes where the decision is represented. It does not make $x$ static or remove the negative case.

Check a batched numerical selection

Check a batched numerical selection · Python experiment

# Experiment — Check a batched numerical selection: Correct values under batching do not imply a conditional branch...
values=jnp.array([-3.,-1.,2.,4.])
assert jnp.allclose(jax.vmap(magnitude)(values),jnp.abs(values))
assert jnp.allclose(jax.vmap(jax.grad(magnitude))(values),jnp.array([-1.,-1.,1.,1.]))

Expected: Magnitudes $[3,1,2,4]$ and slopes $[-1,-1,1,1]$ .

Correct values under batching do not imply a conditional branch remains lazy after every transformation.

Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/03-state/04-branches-with-lax-cond/code/main.py (or inside the Make it yours cell of exercises/state-04.ipynb) and run python3 phases/03-state/04-branches-with-lax-cond/code/main.py to verify your prediction before opening the reference solution.

Write a compiled function that applies $2x$ when $x$ is positive and $x-2$ otherwise. Check positive, negative, and zero inputs.

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.jit(fn) / @jax.jit — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

2. Step-by-step implementation plan:

Define and JIT-compile choose_update(x) so XLA traces and fuses the operations:

Starter code scaffold · copy into your file & fill in the TODOs

# Exercise solution: Write a compiled function that applies 2x when x is positive and x-2...
# Define and JIT-compile `choose_update(x)` so XLA traces and fuses the operations:
@jax.jit
def choose_update(x):
    return ...  # TODO: return computed result
assert jnp.allclose(choose_update(3.), 6.)  # TODO: complete assertion check
assert jnp.allclose(choose_update(-3.), -5.)  # TODO: complete assertion check
assert jnp.allclose(choose_update(0.), -2.)  # TODO: complete assertion check

Reveal the reference solution

Reference solution

# Exercise solution: Write a compiled function that applies 2x when x is positive and x-2...
# Define and JIT-compile `choose_update(x)` so XLA traces and fuses the operations:
@jax.jit
def choose_update(x):
    return jax.lax.cond(x > 0., lambda v: v * 2., lambda v: v - 2., x)
assert jnp.allclose(choose_update(3.), 6.)
assert jnp.allclose(choose_update(-3.), -5.)
assert jnp.allclose(choose_update(0.), -2.)

Match a changed piecewise rule

Practice

Implement $x^2$ for $x>1$ and $2x$ otherwise. Derive outputs at $-1$ , $1$ , $2$ and derivatives away from $1$ .

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.grad(loss_fn)(params, ...) — Transforms a scalar-output function into a function returning the gradient PyTree with the same structure as params.
jax.vmap(fn, in_axes=..., out_axes=...) — Vectorizes a single-example function across a batch axis without writing a Python loop.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

Starter code scaffold · copy into your file & fill in the TODOs

# Match a changed piecewise rule (Practice): The values check the exact boundary and both formulas; the...
def piecewise(v):
    return ...  # TODO: return computed result
assert jnp.allclose(jax.vmap(piecewise)(jnp.array([-1.,1.,2.])),jnp.array([-2.,2.,4.]))  # TODO: complete assertion check
assert jnp.allclose(jax.grad(piecewise)(-1.),2.)  # TODO: complete assertion check
assert jnp.allclose(jax.grad(piecewise)(2.),4.)  # TODO: complete assertion check

Show a hint

At $x=1$ the false branch is selected.

Reveal reference solution and reasoning

Match a changed piecewise rule · Python solution

# Match a changed piecewise rule (Practice): The values check the exact boundary and both formulas; the...
def piecewise(v):
    return jax.lax.cond(v>1.,lambda z:z*z,lambda z:2*z,v)
assert jnp.allclose(jax.vmap(piecewise)(jnp.array([-1.,1.,2.])),jnp.array([-2.,2.,4.]))
assert jnp.allclose(jax.grad(piecewise)(-1.),2.)
assert jnp.allclose(jax.grad(piecewise)(2.),4.)

The values check the exact boundary and both formulas; the derivatives are checked only at smooth points.

Repair incompatible branch shapes

Challenge

Return a scalar in one branch and a length-two vector in the other, reproduce the error, then define a consistent two-value result.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

2. Step-by-step implementation plan:

Run the boundary check and catch the expected exception:

Starter code scaffold · copy into your file & fill in the TODOs

# Repair incompatible branch shapes (Challenge): A branch mismatch is a result-contract error even when the...
# Run the boundary check and catch the expected exception:
try:
    jax.lax.cond(True,lambda x:x,lambda x:jnp.stack([x,x]),jnp.array(2.))
except TypeError:
    print("Expected branch shape mismatch")
else:
    raise AssertionError("Expected result type error")
def pair(v):
    return ...  # TODO: return computed result
assert jnp.allclose(pair(jnp.array(-2.)),jnp.array([2.,4.]))  # TODO: complete assertion check

Show a hint

Decide what each output component means before padding.

Reveal reference solution and reasoning

Repair incompatible branch shapes · Python solution

# Repair incompatible branch shapes (Challenge): A branch mismatch is a result-contract error even when the...
# Run the boundary check and catch the expected exception:
try:
    jax.lax.cond(True,lambda x:x,lambda x:jnp.stack([x,x]),jnp.array(2.))
except TypeError:
    print("Expected branch shape mismatch")
else:
    raise AssertionError("Expected result type error")
def pair(v):
    return jax.lax.cond(v>=0.,lambda x:jnp.stack([x,x*x]),lambda x:jnp.stack([-x,x*x]),v)
assert jnp.allclose(pair(jnp.array(-2.)),jnp.array([2.,4.]))

A branch mismatch is a result-contract error even when the current predicate is known. Here the repaired output explicitly means [magnitude,square].

Diagnose the result

A branch mismatch is a result-contract error even when the current predicate is known. Here the repaired output explicitly means [magnitude,square].

Carry forward

The repair changes where the decision is represented. It does not make $x$ static or remove the negative case.
Correct values under batching do not imply a conditional branch remains lazy after every transformation.

Keep your evidence

Download evidence template

Keep the branch-contract diagram, eager/jit failure and repair, boundary cases, batched values and off-boundary derivatives, and the shape-mismatch diagnosis.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references

Phase 03 Integration Project

Bring this phase together: A deterministic batched simulation

Checkpoint to prove: Replay the same experiment with saved state and keys; explain the update order.

Copy projects/foundation-toolkit/starter/toolkit.py to projects/foundation-toolkit/my_toolkit.py and write your code in projects/foundation-toolkit/my_toolkit.py. Run python3 projects/foundation-toolkit/tests/check.py --implementation projects/foundation-toolkit/my_toolkit.py --stage 4 from the top-level folder to verify stage 4.

Open project workspace · stages 4 Read project guide & figures

Branches with lax.cond

What you will be able to do

The problem

The idea

Trace both possibilities, execute the selected rule

Locate the decision boundary

Trace both branches and check their result types

Choose cond or where by the shape of the question

Inspect boundary behavior and derivatives

Both branches form one magnitude function

Read the figure

Connect it to the computation

1 Prepare the inputs

2 Build the computation

3 Run and check the result

Run the example

Recorded reference execution

Reproduce the Python truth-test failure

Check a batched numerical selection

Make it yours

Match a changed piecewise rule

Repair incompatible branch shapes

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references

Bring this phase together: A deterministic batched simulation