Compute row means while ignoring NaNs

When NaN marks a missing observation, calculate the mean using only the observations that are present:

import torch

x = torch.tensor([1., float("nan"), 3.])
print(x.mean().item())  # nan
print(x.nansum().item())  # 4.0
print(x.nanmean().item())  # 2.0

Ordinary sum and mean propagate NaNs in the values they reduce. nansum excludes NaNs from the sum. nanmean excludes them from both the sum and the count, giving (1 + 3) / 2 = 2 here.

Why replacing missing values with zero changes the answer

Replacing the NaN with zero and then calling mean gives (1 + 0 + 3) / 3. The denominator still counts the missing observation. That answers a different question: the mean after imputing zero, rather than the mean of the available observations.

An entirely missing row has no observations to average. nanmean returns NaN for that row. Returning zero would incorrectly turn missing information into a measured value.

Decide what NaN means in your data

Ignoring NaNs is appropriate when they intentionally mark missing values, such as a sensor dropout. If they came from invalid arithmetic during training, skipping them can hide the underlying problem. Check where they first appear with torch.isnan and inspect the calculation that produced them.

nanmean skips NaNs, not all nonfinite values. Positive and negative infinity are not treated as missing. Also, nan_to_num replaces infinities by default as well as NaNs, so it is not interchangeable with a NaN-aware reduction.

Floating-point numbers have limited range and precision

A floating-point number stores significant digits and an exponent. The exponent limits its range; the number of significant digits limits its precision. Overflow occurs when a result is too large to represent. Very small values can lose precision or round to zero through underflow. Ordinary rounding can also make nearby values indistinguishable.

Subtracting nearly equal values can leave a result dominated by earlier rounding errors. For tiny x, computing log(1 + x) can fail because 1 + x rounds to 1 before the logarithm is taken. The specialized function log1p(x) avoids that loss. Similarly, expm1(x) computes exp(x) - 1 accurately near zero.

Use formulas with safer intermediate values

A numerically stable formula computes the same mathematical quantity while reducing avoidable rounding error or overflow. For finite, nonempty inputs, log-sum-exp can be computed as

$$\log\sum_i e^{z_i}=m+\log\sum_i e^{z_i-m},\quad m=\max_i z_i.$$

After subtracting m, every exponent is at most zero. This avoids exponentiating large positive numbers. For scores [1000, 1000], the result is 1000 + log(2) without ever computing exp(1000).

Compare with a tolerance appropriate to the task

Floating-point addition is not associative: changing the order of a sum can change its rounded result. Parallel reductions and different matrix-multiplication kernels can therefore produce slightly different answers.

Approximate comparisons usually allow an absolute error near zero and a relative error that scales with the reference value. Choose tolerances suited to the problem's scale and dtype. A loose tolerance should not hide a wrong algorithm.

NaN does not compare equal to itself. Decide how invalid values should be handled rather than assuming an ordinary equality test will catch them. Replacing a NaN after an invalid operation may still leave invalid gradients. When differentiating, avoid the invalid operation itself where possible.