Skip to content

← JAX Foundations step 22 of 25

Medium Primitives

Multi-Condition where

Jump to the task

Why this matters

jnp.where(condition, a, b) chooses a value at each position: from a when the condition is true, and from b otherwise. A scalar choice such as -1.0 broadcasts across the array.

Nesting a second selection lets you express three cases without asking Python to decide a separate branch for each array element.

Worked mini-example

import jax.numpy as jnp

x = jnp.array([-3.0, -1.0, 0.0, 4.0, 5.0])
lower, upper = -1.0, 4.0
out = jnp.where(x < lower, -1.0, jnp.where(x > upper, 1.0, 0.0))
# out = [-1, 0, 0, 0, 1]

Values on either boundary belong to the middle interval. The comparisons are strict: less than the lower bound and greater than the upper bound.

For a longer list of conditions, jnp.select can be easier to read. The first true condition wins when conditions overlap:

same = jnp.select([x < lower, x > upper], [-1.0, 1.0], default=0.0)
# same contains the same five values as out.

Common pitfalls

  • Python if cannot decide a dynamic JIT-traced comparison. It raises TracerBoolConversionError. Array selection with where lets different elements take different cases.
  • The condition and both choices must be broadcast-compatible. Scalar choices work for every position here.
  • where is selection, not lazy Python branching. Do not assume an unselected expression is skipped. Invalid operations in a candidate expression can also cause NaN gradients. If using something such as log, make its input valid before applying it. This exercise only selects constants.
  • Use < and >, not <= or >=: the boundaries should map to zero.

Problem

Implement piecewise_classify(x, lower, upper). The input x is a one-dimensional finite floating-point JAX array. The finite scalar bounds satisfy lower <= upper.

Return an array with the same shape, containing:

  • -1.0 for values below lower.
  • 1.0 for values above upper.
  • 0.0 for values in the inclusive interval between them.

If both bounds are equal, only that exact value belongs to the middle interval. An empty input produces an empty output.