Skip to content

← JAX Foundations step 23 of 25

Easy Primitives

NaN-Safe Mean (with mask)

Jump to the task

Average only the values you want to include

A padded sequence or a table with missing observations may contain positions that should not count toward a mean. Exclude those positions from both the sum and the count. Replacing them with zero and then taking an ordinary mean would still include them in the denominator.

In this exercise, a separate mask marks excluded positions. It is the source of truth: do not infer validity from the values in x. Standard JSON cannot encode NaN, so the fixtures use ordinary numbers as placeholders and supply the mask separately.

Worked example

import jax.numpy as jnp

x = jnp.array([2.0, 999.0, 6.0])
nan_mask = jnp.array([0.0, 1.0, 0.0])
valid = nan_mask < 0.5
cleaned = jnp.where(valid, x, 0.0)  # [2.0, 0.0, 6.0]
count = jnp.sum(valid)            # 2
mean = jnp.sum(cleaned) / jnp.maximum(count, 1)  # 4.0

The placeholder 999 contributes neither to the numerator nor to the denominator. Dividing the cleaned sum by the array length would give 8 / 3, which is not the mean of the included observations.

When nothing is included, the cleaned sum is zero. Using a denominator of at least one returns zero without dividing by zero. This is the exercise's chosen fallback, not the mathematical mean of an empty set.

Problem

Implement safe_mean_with_nan(x, nan_mask) to return the mean of values whose mask is less than 0.5. A mask value equal to or greater than 0.5 excludes its position. Return a scalar JAX floating-point value.

  • x and nan_mask are one-dimensional float32 arrays of the same shape.
  • Mask entries are finite numbers. The usual convention is 0 for included and 1 for excluded, but apply the specified threshold.
  • If all entries are excluded, or the arrays are empty, return 0.0.

For example, [2, 4, 6] with mask [0, 1, 0] gives 4.0. [5] with mask [1] gives 0.0.

Missing values and invalid calculations

If real data uses NaN itself to mark missing entries, ~jnp.isnan(x) can build a validity mask. That is a different input convention from this exercise. Here, an unmasked NaN remains part of the calculation and makes the result NaN; infinity is not automatically treated as missing either. jnp.nanmean ignores NaNs but returns NaN for an all-NaN slice, so its default empty-result policy differs from the one requested here.

Use jnp.where to select zero at excluded positions. Multiplying by a zero mask does not remove NaN, because NaN * 0 is still NaN. Keeping the array shape fixed also avoids the dynamic output size of boolean indexing under JIT.

Selection can remove an already-present NaN from this forward sum, but it does not make an earlier invalid operation safe to differentiate. For example, masking the output of a logarithm does not guarantee finite gradients through that logarithm. Make its inputs valid before evaluating it when you need gradients through the computation.