Skip to content

← Optax step 20 of 25

Medium Primitives

Simulated NaN Cleanup Before SGD

Jump to the task

Replace selected gradient entries with zero, then take one plain SGD step. The mask stands in for NaN locations because the exercise's JSON inputs contain ordinary finite numbers rather than NaN values.

The mask is authoritative

A mask value of 1.0 means ignore that gradient entry. A value of 0.0 means use it, regardless of its magnitude. For example, with params = [1, 2, 3], grads = [0.1, 999, 0.3], mask = [0, 1, 0] and lr = 0.5, use cleaned gradients [0.1, 0, 0.3] and return [0.95, 2.0, 2.85].

The middle parameter stays at 2.0; it is the gradient that becomes zero. The value 999 is just a placeholder, not a special marker. Checking for NaNs in these finite inputs would not identify the selected positions.

Relationship to zero_nans

optax.zero_nans() replaces actual NaN entries in incoming updates with zero. Here, implement the corresponding selection with jnp.where(mask >= 0.5, 0.0, grads), then pass the cleaned array to optax.sgd(lr).

Replacing NaNs does not repair the calculation that produced them or ensure valid training. It also does not remove infinities. Norm clipping is not a reliable way to repair infinite gradients: scaling an infinite entry can itself produce NaN. This exercise only defines the finite, mask-based update.

Inputs

Implement safe_step_with_nan_zeroing(params, grads, mask, lr). Use plain SGD without momentum or weight decay and perform one update.

  • params and grads: finite 1-D floating-point JAX arrays with the same shape.
  • mask: a JAX array of that shape containing only 0.0 and 1.0. A value of 1.0 selects a gradient to replace with zero.
  • lr: a scalar learning rate.

Output

Return the updated JAX array with the original shape. Leave all input arrays unchanged. Selected parameter values remain at their original values; unselected positions receive the usual SGD update.