Skip to content

← JAX Vectorization & Control Flow step 15 of 25

Easy Primitives

Binary Classification via lax.cond

Jump to the task

Choose a label from a scalar condition

This exercise uses a simple threshold rule: add the input values, compare that total with a threshold, and choose one of two labels. It illustrates branching, not training a classifier.

Under jax.jit, a Python if statement cannot choose a branch from a condition whose value depends on runtime input. Such a condition has no concrete Python boolean during tracing. lax.cond represents that decision inside the compiled program instead.

import jax
from jax import lax

@jax.jit
def positive_label(x):
    return lax.cond(x > 0, lambda: 1.0, lambda: -1.0)

positive_label(3.0)   # 1.0
positive_label(-2.0)  # -1.0
positive_label(0.0)   # -1.0

Zero takes the false branch. This label rule is different from jnp.sign, which returns zero at zero.

Branch functions must agree on their outputs

Here the branch functions take no arguments because they only return constants. More generally, lax.cond(pred, true_fn, false_fn, operand) passes operand to whichever function is selected. The predicate must be scalar; use array selection or batching for elementwise decisions.

Both branches must return the same structure, array shapes and dtypes. Returning integer 1 from one branch and floating-point -1.0 from the other can fail during tracing, even without an outer JIT. Using 1.0 and -1.0 keeps this example's outputs consistent.

Separate tracing from execution

Both branch functions are traced, so both must contain valid JAX operations. Ordinary Python prints or mutations can happen during tracing; they are not reliable evidence that a runtime branch ran. Keep branch functions pure.

A scalar cond expresses conditional execution, subject to compiler transformations. It is not generally lowered to elementwise select. When vmap batches the predicate, however, cond becomes selection and both branch computations can run. Do not assume cond always avoids the unused computation under every transformation.

jnp.where instead selects between supplied values or arrays. Its branch expressions are evaluated to form those arguments, though the compiler may remove unnecessary work. It is useful for elementwise selection, but neither API promises a speedup for this tiny label rule.

Problem

Implement conditional_classify(x, threshold) using lax.cond. The input x is a one-dimensional real floating-point JAX array and threshold is a real floating-point scalar. Return scalar 1.0 when jnp.sum(x) > threshold, otherwise scalar -1.0.

Compare the sum once; do not classify each element separately or replace the sum with a mean. Equality takes the false branch. An empty input has sum zero. For x=[1, 2] and threshold=3, return -1.0; for the same x and threshold=2, return 1.0.

Pass branch functions to cond, not their evaluated results. Output tests verify labels but cannot enforce the choice of conditional API.