Skip to content

← JAX Foundations step 25 of 25

Easy Primitives

Clip and Extrema

Jump to the task

Limit values, then apply a baseline

Clipping keeps each value within a chosen interval. For bounds lo and hi, values below lo become lo, values above hi become hi, and values already inside the interval stay unchanged.

This exercise combines three elementwise operations: clip, subtract a baseline, then replace negative results with zero. Each input value produces one output value; no reduction across the array is needed.

Worked example

import jax.numpy as jnp

x = jnp.array([-2.0, 0.5, 3.0])
clipped = jnp.clip(x, 0.0, 1.0)      # [0.0, 0.5, 1.0]
shifted = clipped - 0.3             # approximately [-0.3, 0.2, 0.7]
result = jnp.maximum(shifted, 0.0)   # approximately [0.0, 0.2, 0.7]

The original 3 first becomes 1, then 0.7. Subtracting before clipping would instead produce 1 at that position, so order matters. These operations return array values without changing x.

Choose the right maximum

jnp.maximum(a, b) compares corresponding elements. A scalar argument broadcasts across an array, so jnp.maximum(x, 0) sets a lower bound of zero at every position. jnp.minimum similarly applies an upper bound.

jnp.max(x) is a reduction: without an axis argument, it returns the single largest value. It cannot replace jnp.maximum(x, 0) here. Flooring at zero means applying a lower bound, not rounding with jnp.floor.

Problem

Implement clip_and_subtract_baseline(x, lo, hi, baseline):

  1. Clip x to the inclusive interval [lo, hi].
  2. Subtract baseline from each clipped value.
  3. Replace negative results with zero.

x is a one-dimensional float32 JAX array. The bounds and baseline are finite scalars, with lo <= hi. Return an array with the same shape as x. Empty input gives empty output. The final values need not lie in [lo, hi], since the baseline and lower bound are applied after clipping.

For example, x = [-1, 3], lo = 0, hi = 2, and baseline = 1.5 give [0, 0.5]. A negative baseline increases the clipped values before the final lower bound is applied.

Limits of clipping

Clipping does not diagnose or repair invalid calculations. NaNs remain NaN through this pipeline. If a NaN is unexpected, investigate its source; if it represents missing data, choose an explicit missing-data policy. Also, elementwise clipping is different from clipping a vector's norm, which rescales the vector according to its total length.