Skip to content

← JAX Vectorization & Control Flow step 17 of 25

Hard Primitives

vmap over lax.cond

Jump to the task

Give each row its own label

Start with a rule for one row: add its values, compare the sum with its threshold, and return 1.0 or -1.0. This is a fixed rule, not a trained classifier. vmap applies that rule to every row, pairing each row with the threshold at the same position.

With a batched predicate, vmap transforms lax.cond into selection. Both branches are traced, and computations for both branches may execute. Do not rely on this composition to skip expensive work or to control Python side effects. Keep branch functions pure.

Worked mini-example

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

def classify(row, t):
    return lax.cond(
        row.sum() > t,
        lambda: 1.0,
        lambda: -1.0,
    )

rows = jnp.array([
    [1.0, 2.0],
    [-1.0, -2.0],
    [5.0, 5.0],
])
thresholds = jnp.array([
    2.0, 0.0, 5.0
])
batched = jax.vmap(
    classify,
    in_axes=(0, 0),
)
result = batched(
    rows, thresholds
)
# → [1.0, -1.0, 1.0]

in_axes=(0, 0) tells vmap to map over the leading axis of both arguments simultaneously.

Common pitfalls

  • Passing a result instead of a function. vmap needs a callable, not the array returned by lax.cond(...). A helper exposes the row and threshold as array arguments while keeping branch functions fixed.
  • Reducing the wrong dimension. Inside classify, row is one length-two vector and t is one scalar threshold. row.sum() reduces that row, not the entire batch.
  • Shared versus per-row thresholds. The default in_axes=0 maps both array arguments and works here, just like (0, 0). Use None only for an argument that should be shared across all examples.

In this example, each length-two row is paired with one scalar threshold. Its scalar label ([]) is stacked with the others to give [3]. The diagram shows shapes, not which branch computations execute.

Loading visualization…

Problem

Implement vmap_classify_batch(x_batch, thresholds) that classifies each row of x_batch independently: output i is 1.0 if sum(x_batch[i]) > thresholds[i], else -1.0.

  • x_batch: 2-D JAX array of shape (N, d).
  • thresholds: 1-D JAX array of shape (N,).
  • Returns: floating-point labels with shape (N,).

Equality returns -1.0. For example, row [1, 2] with threshold 3 returns -1.0. An empty batch returns shape (0,); a row with no elements has sum zero. Inputs are real floating-point arrays.

Use a vmapped helper containing lax.cond for practice. Output tests check the labels, but cannot distinguish this from equivalent code.