Skip to content

← JAX Profiling & Debugging step 8 of 12

Medium Primitives

Fix: Python if → jnp.where

Jump to the task

Make one decision per element

A threshold mask records which entries exceed a cutoff. For x = [-1, 0.5, 2, 5] and threshold = 1, the comparison x > threshold is [False, False, True, True]. This task turns that mask into the numeric array [0, 0, 1, 1]. Equality belongs to the zero case.

import jax
import jax.numpy as jnp

@jax.jit
def threshold_mask(x, threshold):
    return jnp.where(x > threshold, 1.0, 0.0)

result = threshold_mask(jnp.array([-1., 0.5, 2., 5.]), 1.0)
# [0., 0., 1., 1.]

Both the array values and threshold can change between calls. Their values do not need to be static because the comparison and selection are represented in the compiled calculation.

A Python if makes a different kind of choice

def first_element_decides(x, threshold):
    if x[0] > threshold:
        return jnp.ones_like(x)
    return jnp.zeros_like(x)

With concrete inputs, this decides the entire output from x[0]. For the example above it returns all zeros, not [0, 0, 1, 1]. It is not an equivalent implementation of the elementwise task.

Under ordinary JIT with dynamic inputs, the Python if also raises TracerBoolConversionError: Python needs a boolean while tracing, but the comparison's value is not yet known. A static condition, such as one based on an ordinary known shape, is a different case and can work in Python control flow.

If the intended behavior really were one whole-array choice, use lax.cond with a scalar predicate and two branch functions. Here the intended behavior is per element, so use the array mask with where.

Selection is not lazy Python branching

The three-argument where takes a condition and two value expressions. Do not rely on it to avoid forming an invalid expression: it is a selection operation, not a short-circuiting if. Compiler optimization can alter the actual executed work, so avoid unconditional claims about how many kernels or branch operations must run.

For example, hiding log(0) behind where can give a finite selected value but a NaN gradient. A numerically safe expression must protect the input to the problematic operation itself. This exercise selects between finite constants, so that issue does not arise here.

lax.cond traces both branch functions and requires compatible results. An ordinary scalar cond represents one runtime branch, but batching predicates with vmap can turn it into selection. It is not a universal way to avoid evaluating an unsafe expression in every transformed case.

Problem

Implement safe_threshold(x, threshold) using jnp.where. Return 1.0 for each entry strictly greater than threshold and 0.0 otherwise.

X is a one-dimensional finite float32 JAX array. Threshold is a finite Python float or scalar float32 JAX array. Return a float32 JAX array with the same shape. Empty input stays empty. Values equal to the threshold must map to zero. Do not filter entries out or let x[0] decide the whole output.

The function must work under jax.jit with a dynamic threshold. Output tests check numerical behavior; inspect the requested use of where separately.