Skip to content

← JAX Vectorization & Control Flow step 20 of 25

Hard Primitives

vmap-cond vs where: cost tradeoff

Jump to the task

Two ways to choose a result for each row

This exercise expresses the same row-wise arithmetic with vmap plus lax.cond, and with jnp.where. They produce the same values here, but neither form guarantees that an unselected branch's work is skipped.

A scalar lax.cond represents a conditional. When its predicate is batched by vmap, JAX converts it to selection. Both branches are traced, and the selected result may require computing both candidate arrays. jnp.where also selects between candidate values. Compiler optimizations can remove unnecessary work, so measure the compiled workload rather than assuming a speed advantage from either spelling.

Branch functions should be pure. Batched cond is not a way to make per-row Python side effects conditional.

Worked mini-example

import jax
import jax.numpy as jnp
from jax import lax

rows = jnp.array([
    [1.0, 2.0, 3.0],
    [-1.0, -2.0, -3.0],
])
t = 2.0

def per_row(row):
    total = row.sum()
    above = total > t
    return lax.cond(
        above,
        lambda: row * 2,
        lambda: row + 1,
    )

via_cond = jax.vmap(
    per_row
)(rows)
above = rows.sum(axis=1)
above = above > t
via_where = jnp.where(
    above[:, None],
    rows * 2,
    rows + 1,
)
result = jnp.stack([
    via_cond, via_where
])
# Both slices:
# [[2, 4, 6],
#  [0, -1, -2]]
# Shape: (2, 2, 3).

Broadcast a row decision across its columns

A condition with shape (N,) aligns with the last dimension, not the row dimension. Adding [:, None] gives shape (N, 1), so each row's decision applies to all d columns. With a square array this mistake can produce a plausible but wrong result instead of an error.

For rows [[2, 3], [-2, -3]] and threshold 0, the correct result is [[4, 6], [-1, -2]]. A condition left at shape (2,) instead selects by column and gives [[4, 4], [-4, -2]].

Common pitfalls

  • Forgetting to broadcast the condition in jnp.where. If x_batch has shape (N, d) and your per-row condition has shape (N,), you must expand it: cond[:, None]. Without this, jnp.where will error or silently misbroadcast.
  • Assuming short-circuit evaluation. Both forms select the same arithmetic results here. Do not use batched cond to protect an unsafe computation or assume it avoids an expensive branch.
  • vmap of cond on 2-D input. per_row_cond receives a 1-D row; the lambda: row * 2.0 must close over that per-row row, not the full batch.

The diagram shows only the vmap stage: a length-three row becomes a length-three result, and two rows stack into [2, 3]. The scalar threshold is captured by per_row. Stacking this result with the where result adds the final leading axis, producing [2, 2, 3].

Loading visualization…

Problem

Implement vmap_branch_two_ways(x_batch, threshold) that applies the branch logic row * 2 if sum(row) > threshold else row + 1 per row, via both vmap(lax.cond) and jnp.where.

  • x_batch: 2-D JAX array of shape (N, d).
  • threshold: scalar.
  • Returns: 3-D array of shape (2, N, d): [result_via_cond, result_via_where]. Both slices are numerically identical.

Equality takes the row + 1 branch. Inputs are real floating-point arrays and a real scalar threshold; preserve the array dtype. Empty batches or rows preserve their dimensions in (2, N, d). The threshold is shared by all rows, unlike a per-row threshold array.

This exercise checks values and shapes, not execution cost. Output tests cannot verify that you used two different implementations. For a separate timing experiment, compile and warm up each method, then wait for each result with block_until_ready() when measuring it.