Introduction

NaN values can silently destabilize TensorFlow training and inference if you do not handle them early. They often appear after division by zero, invalid logarithms, or overflow in preprocessing. Replacing NaN safely requires both a correct tensor operation and a clear policy for replacement values.

Detect NaN Values in Tensors

Before replacing values, inspect where NaN appears. TensorFlow offers tf.math.is_nan for element wise detection.

1import tensorflow as tf
2
3x = tf.constant([1.0, float("nan"), 3.5, float("nan")], dtype=tf.float32)
4mask = tf.math.is_nan(x)
5print(mask.numpy())

This mask is useful for debugging and for calculating NaN ratio metrics in data pipelines.

Replace NaN With tf.where

The most common replacement strategy uses tf.where with the NaN mask.

fill_value = tf.constant(0.0, dtype=tf.float32)
clean = tf.where(tf.math.is_nan(x), fill_value, x)
print(clean.numpy())

This keeps original non NaN values and substitutes only problematic entries.

For matrix tensors, the same pattern works with broadcasting.

m = tf.constant([[1.0, float("nan")], [float("nan"), 4.0]], dtype=tf.float32)
clean_m = tf.where(tf.math.is_nan(m), tf.zeros_like(m), m)
print(clean_m.numpy())

Use tf.experimental.numpy.nan_to_num Style Replacement

If you need broader numeric cleanup including positive or negative infinity, tf.where chains can work, but a helper approach is often clearer.

1def sanitize_tensor(t):
2    t = tf.where(tf.math.is_nan(t), tf.zeros_like(t), t)
3    t = tf.where(tf.math.is_inf(t), tf.zeros_like(t), t)
4    return t
5
6z = tf.constant([1.0, float("inf"), float("nan"), -2.0], dtype=tf.float32)
7print(sanitize_tensor(z).numpy())

This gives one reusable sanitization function for preprocessing and model serving.

Choose Replacement Values Carefully

Zero is common, but not always correct. In standardized features, mean imputation may be better. In log scaled features, a small epsilon might preserve scale assumptions better than zero.

For sequence models, replacing with zero can introduce a token like pattern in numeric space. In that case, keep a companion mask feature so the model can learn which values were originally missing.

Add Numeric Guards During Training

Replacing NaN in input tensors is only one side of stability. During training, you should also guard gradient updates and monitor loss values for sudden divergence.

1optimizer = tf.keras.optimizers.Adam(1e-3)
2
3for step, (features, labels) in enumerate(train_ds):
4    features = sanitize_tensor(features)
5    with tf.GradientTape() as tape:
6        preds = model(features, training=True)
7        loss = loss_fn(labels, preds)
8
9    grads = tape.gradient(loss, model.trainable_variables)
10    grads = [tf.where(tf.math.is_finite(g), g, tf.zeros_like(g)) for g in grads]
11    optimizer.apply_gradients(zip(grads, model.trainable_variables))
12
13    if step % 100 == 0:
14        tf.debugging.assert_all_finite(loss, "Loss became non finite")

This pattern keeps occasional numeric spikes from poisoning the entire optimization run.

Common Pitfalls

• Replacing NaN without investigating source: hidden numeric bugs remain.
• Ignoring infinity values: model still receives invalid numbers.
• Mixing dtypes during replacement: implicit casts may reduce precision.
• Applying sanitization after batching only: bad values may already affect augmentations.
• Hardcoding one fill value for all features: statistical meaning can be lost.

Summary

• Detect NaN with tf.math.is_nan before replacement.
• Use tf.where for explicit and controllable element wise cleanup.
• Sanitize infinity values along with NaN for robust pipelines.
• Pick replacement values based on feature semantics, not convenience.
• Track missingness where needed so models retain useful context.
• Add finite value assertions during training loops so numeric instability is detected immediately instead of propagating silently through checkpoints.
• Track the percentage of replaced values over time to catch upstream data quality regressions early.
• Document fill value rationale so future maintainers preserve statistical intent.