Both branches, one index

Put a print inside each branch function and trace the function once. Both prints fire. A jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → for each branch has to exist before the primitive can be built, so your Python runs once per branch, at trace time, whatever the predicate later turns out to be.

The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → then shows what became of the predicate. gt produced a bool[], and the very next equation converts it to i32[]. cond never receives a boolean at all; it receives an index and a tuple of branches.

Read the branch order twice, because it is the reverse of how the call is written. div sits in the first slot and mul in the second, and div was the false branch. One line in the source says why: cond_p.bind(index, *consts, *ops, branches=(false_jaxpr, true_jaxpr)). Slot 0 is false, slot 1 is true, so a two-way cond is a two-entry switch that was handed a converted boolean.

import jax
import jax.numpy as jnp

def flip(x):
    def hot(v):
        print("tracing hot")
        return v * 2.0
    def cold(v):
        print("tracing cold")
        return v / 2.0
    return jax.lax.cond(x.sum() > 0, hot, cold, x)

print(jax.make_jaxpr(flip)(jnp.ones(3)))

# tracing hot
# tracing cold
# { lambda ; a:f32[3]. let
#     b:f32[] = reduce_sum[axes=(0,)] a
#     c:bool[] = gt b 0.0
#     d:i32[] = convert_element_type[new_dtype=int32 weak_type=False] c
#     e:f32[3] = cond[
#       branches=(
#         { lambda ; f:f32[3]. let g:f32[3] = div f 2.0 in (g,) }
#         { lambda ; h:f32[3]. let i:f32[3] = mul h 2.0 in (i,) }
#       )
#     ] d a
#   in (e,) }

Hand cond a Python True and nothing collapses. Both branches are still in the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → and the index is the literal 1. The primitive was built the same way it always is; only the value feeding it is known early.

That is worth holding against the chapter above, where the branching that really does resolve during tracing is a Python if on a static value. The two cases look similar in the source and land in different places, and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → is where you tell them apart: an if leaves no cond equation behind at all, and a cond leaves one even when its predicate is a constant.

import jax
import jax.numpy as jnp

def always(x):
    return jax.lax.cond(True, lambda v: v * 2.0, lambda v: v / 2.0, x)

print(jax.make_jaxpr(always)(jnp.ones(3)))

# { lambda ; a:f32[3]. let
#     b:f32[3] = cond[
#       branches=(
#         { lambda ; c:f32[3]. let d:f32[3] = div c 2.0 in (d,) }
#         { lambda ; e:f32[3]. let f:f32[3] = mul e 2.0 in (f,) }
#       )
#     ] 1 a
#   in (b,) }

jax.lax.switch is the n-way form, and its docstring states the semantics as three lines of Python you can hold in your head. Clamp the index into range, then apply that branch. The clamp is not a footnote: it is an equation in the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, sitting right before the cond.

So an index of 7 against three branches runs branch 2, and an index of -5 runs branch 0. No error, no warning, and a result that looks like a legitimate answer. If your index is computed from data, the failure mode for an off-by-one is a plausible number rather than a raise.

The rows below came from one jitted function called five times. Only the index changed between calls, and the same executable served all five, because the index is a traced value like any other.

  """Apply exactly one of the ``branches`` given by ``index``.

  If ``index`` is out of bounds, it is clamped to within bounds.

  Has the semantics of the following Python::

    def switch(index, branches, *operands):
      index = clamp(0, index, len(branches) - 1)
      return branches[index](*operands)

# >>> three = [lambda v: v + 1.0, lambda v: v + 2.0, lambda v: v + 3.0]
# >>> def pick(i, x): return jax.lax.switch(i, three, x)
# >>> print(jax.make_jaxpr(pick)(jnp.int32(0), jnp.ones(2)))
# { lambda ; a:i32[] b:f32[2]. let
#     c:i32[] = clamp 0 a 2
#     d:f32[2] = cond[
#       branches=(
#         { lambda ; e:f32[2]. let f:f32[2] = add e 1.0 in (f,) }
#         { lambda ; g:f32[2]. let h:f32[2] = add g 2.0 in (h,) }
#         { lambda ; i:f32[2]. let j:f32[2] = add i 3.0 in (j,) }
#       )
#     ] c b
#   in (d,) }

The branches meet at one output type, and the check is exact rather than sympathetic. Different shapes raise. Different dtypes raise. Different pytree structures raise, and that one is reported separately, naming both treedefs instead of comparing avals.

One disagreement is allowed through, and it is the one people expect to break. A branch returning the Python float 1.0 against a branch returning an f32[] array is fine, because the weak scalar promotes to meet the array. Weakness is a property of the type here, not a special case in cond.

The habit worth building is to read the two avals in the message before rereading your code. DIFFERENT ShapedArray(float32[3]) vs. ShapedArray(float32[2]) tells you which axis moved, and it is usually a slice or a reduction inside one branch that you did not mirror in the other.

import jax
import jax.numpy as jnp

x = jnp.ones(3)
pred = x.sum() > 0

def show(true_fn, false_fn):
    try:
        jax.lax.cond(pred, true_fn, false_fn, x)
        print("accepted")
    except TypeError as err:
        print(err)

show(lambda v: v, lambda v: v[:2])                    # shapes disagree
show(lambda v: v, lambda v: v.astype(jnp.float16))    # dtypes disagree
show(lambda v: (v, v), lambda v: v)                   # structures disagree
show(lambda v: v.sum(), lambda v: 1.0)                # a weak python float

# true_fun and false_fun output must have identical types, got
# DIFFERENT ShapedArray(float32[3]) vs. ShapedArray(float32[2]).
# true_fun and false_fun output must have identical types, got
# DIFFERENT ShapedArray(float32[3]) vs. ShapedArray(float16[3]).
# true_fun and false_fun output must have same type structure, got PyTreeDef((*, *)) and PyTreeDef(*).
# accepted

Close over a value in one branch only and watch what the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → does with it. Both branch signatures grow, both get the value, and the one that has no use for it marks its binder with a trailing underscore. The operand list carries it twice, once for each branch.

The comment above the function that does this states the reason in three sentences: the staged jaxprsThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → are the branches of one primitive, so their input signatures have to match, and the fix is to make each accept all the constants and drop the ones it does not need.

This is where a closure over something large stops being free. Two branches closing over two different big arrays produce one primitive whose operand list holds both, whichever branch runs. Passing the value as an operand instead of closing over it changes nothing about that, but it does make the cost visible at the call site.

  # When staging the branches of a conditional into jaxprs, constants are
  # extracted from each branch and converted to jaxpr arguments. To use the
  # staged jaxprs as the branches to a conditional *primitive*, we need for
  # their (input) signatures to match. This function "joins" the staged jaxprs:
  # for each one, it makes another that accepts *all* constants, but only uses
  # those that it needs (dropping the rest).

# >>> def f(x, w):
# ...     return jax.lax.cond(x.sum() > 0, lambda v: v * w, lambda v: v - w, x)
# >>> print(jax.make_jaxpr(f)(jnp.ones(3), jnp.full((3,), 2.0)))
# { lambda ; a:f32[3] b:f32[3]. let
#     c:f32[] = reduce_sum[axes=(0,)] a
#     d:bool[] = gt c 0.0
#     e:i32[] = convert_element_type[new_dtype=int32 weak_type=False] d
#     f:f32[3] = cond[
#       branches=(
#         { lambda ; g_:f32[3] h:f32[3] i:f32[3]. let
#             j:f32[3] = convert_element_type[new_dtype=float32 weak_type=False] h
#             k:f32[3] = sub i j
#           in (k,) }
#         { lambda ; l:f32[3] m_:f32[3] n:f32[3]. let
#             o:f32[3] = convert_element_type[new_dtype=float32 weak_type=False] l
#             p:f32[3] = mul n o
#           in (p,) }
#       )
#     ] e b b a
#   in (f,) }

Give one branch five chained 512 by 512 matmuls and the other an add, then time the same compiled function under both predicates. The heavy predicate costs milliseconds and the light one costs almost nothing, so on this backend the executable really is skipping the branch it did not select.

Batch that same function and the timing stops depending on the predicate. Two lanes with both predicates false still pay for the matmuls, because the batched form has to produce both results before it can select between them per lane. Chapter 6 states that cost; the stopwatch is here so you can see how large it gets when one branch is expensive.

Which gives a rule for where to put a cond in a batched program. A branch that guards a cheap correction costs little either way. A branch that guards the expensive half of your model is worth hoisting out of the vmap, so the predicate applies to the whole batch instead of per element.

import time
import jax
import jax.numpy as jnp

A = jnp.ones((512, 512))
heavy = lambda v: (A @ A @ A @ A @ A).sum() + v
light = lambda v: v + 1.0
one = jax.jit(lambda p, x: jax.lax.cond(p, heavy, light, x))
many = jax.jit(jax.vmap(lambda p, x: jax.lax.cond(p, heavy, light, x)))

def best(fn, *a):
    fn(*a).block_until_ready()
    ts = []
    for _ in range(50):
        t0 = time.perf_counter()
        fn(*a).block_until_ready()
        ts.append(time.perf_counter() - t0)
    return 1e3 * min(ts)

x, xs = jnp.float32(1.0), jnp.ones(2)
print(f"cond, predicate True   {best(one, jnp.bool_(True), x):6.2f} ms")
print(f"cond, predicate False  {best(one, jnp.bool_(False), x):6.2f} ms")
print(f"vmapped, both False    {best(many, jnp.array([False, False]), xs):6.2f} ms")
print(f"vmapped, both True     {best(many, jnp.array([True, True]), xs):6.2f} ms")

Guard a 1 / x with jnp.where and differentiate it at zero, and back comes nan. The masked-out value was still computed, its derivative was still taken, and multiplying an infinite derivative by a zero cotangent produces the indeterminate form. This is the classic reason people distrust masking.

Write the same guard as a cond, batch it, differentiate it, and the answer is 0. That is not because the batched form skipped anything. It computed both branches, the same way the timings above showed.

The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → says where the protection comes from. Before each branch runs, its operand goes through a select_n against stop_gradient of itself: the lanes that branch owns receive the real value, and every other lane receives a gradient-blocked copy. So the 1 / x still evaluates at zero, and no cotangent from it can reach the input.

Batching a cond costs you both branches. It does not cost you the gradient.

import jax
import jax.numpy as jnp

masked = lambda x: jnp.where(x > 0, 1.0 / x, 0.0)
branched = lambda x: jax.lax.cond(x > 0, lambda v: 1.0 / v, lambda v: 0.0 * v, x)

xs = jnp.array([0.0, 2.0])
print(jax.vmap(masked)(xs), jax.vmap(branched)(xs))
print(jax.vmap(jax.grad(masked))(xs))
print(jax.vmap(jax.grad(branched))(xs))
print(jax.make_jaxpr(jax.vmap(branched))(xs))

# [0.  0.5] [0.  0.5]
# [  nan -0.25]
# [ 0.   -0.25]
# { lambda ; a:f32[2]. let
#     b:bool[2] = gt a 0.0
#     c:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b
#     d:bool[2] = eq c 0
#     e:f32[2] = stop_gradient a
#     f:f32[2] = select_n d e a
#     g:f32[2] = mul 0.0 f
#     h:bool[2] = eq c 1
#     i:f32[2] = stop_gradient a
#     j:f32[2] = select_n h i a
#     k:f32[2] = div 1.0 j
#     l:f32[2] = select_n c g k
#   in (l,) }