TensorFlow
NaN Handling
Data Cleaning
Machine Learning
Python

Tensorflow How to convert NaNs to a number?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

If you want to replace NaN values in a TensorFlow tensor with a real number such as 0.0, the usual solution is tf.where combined with tf.math.is_nan. That gives you explicit control over which values are replaced and what replacement value should be used.

This matters because NaN values propagate through arithmetic and can quickly poison training or inference results. Replacing them early in the input pipeline or right before sensitive computations is often the safest approach.

Replace NaN with tf.where

The direct pattern looks like this:

python
1import tensorflow as tf
2
3x = tf.constant([1.0, float('nan'), 3.5, float('nan')], dtype=tf.float32)
4clean = tf.where(tf.math.is_nan(x), 0.0, x)
5
6print(clean.numpy())

This returns a tensor where every NaN has been replaced by 0.0 while non-NaN values are left unchanged.

You can replace with any number you want:

python
clean = tf.where(tf.math.is_nan(x), -1.0, x)

That is often useful when a specific sentinel value or default fill is required.

Match the Tensor Shape and Type

For scalar replacement values such as 0.0, TensorFlow broadcasts automatically. If you need a tensor-shaped replacement, provide one with a compatible shape and dtype:

python
replacement = tf.fill(tf.shape(x), 5.0)
clean = tf.where(tf.math.is_nan(x), replacement, x)

This is helpful when the replacement value is computed dynamically rather than hard-coded.

Clean Matrices and Higher-Dimensional Tensors

The same pattern works for matrices and batches:

python
1x = tf.constant([
2    [1.0, float('nan')],
3    [float('nan'), 4.0]
4], dtype=tf.float32)
5
6clean = tf.where(tf.math.is_nan(x), tf.zeros_like(x), x)
7print(clean.numpy())

Using tf.zeros_like is a convenient way to replace all NaN values with zeros while keeping the same shape and dtype.

Consider Why the NaN Values Exist

Replacing NaN values is useful, but it is also worth asking where they came from. Common causes include:

  • division by zero
  • invalid logarithms
  • exploding gradients
  • missing or badly parsed input data

If the NaN values indicate a modeling bug, replacement may only hide the symptom. In data-cleaning pipelines, replacement is often correct. In model training, it can be a sign that something upstream needs attention.

That is why many teams do both: they replace NaN values at the data boundary for robustness and they also alert or inspect metrics so the root cause is not forgotten. The replacement step should be deliberate, not just automatic cleanup with no follow-up.

Apply It in the Input Pipeline

If the problem comes from training data, cleaning it in the input pipeline keeps the model code simpler:

python
1import tensorflow as tf
2
3def clean_batch(features):
4    return tf.where(tf.math.is_nan(features), 0.0, features)
5
6dataset = tf.data.Dataset.from_tensor_slices(
7    tf.constant([1.0, float('nan'), 2.0], dtype=tf.float32)
8)
9dataset = dataset.map(clean_batch)

That ensures the model never sees the NaN values at all.

Common Pitfalls

  • Replacing NaN values without investigating whether they come from a deeper numerical bug.
  • Forgetting that the replacement value must have a compatible dtype with the tensor.
  • Cleaning only one tensor in a pipeline while other features or labels still contain NaN values.
  • Assuming replacement is always harmless during training when it may distort the data distribution.

Summary

  • Use tf.where(tf.math.is_nan(x), replacement, x) to replace NaN values in TensorFlow.
  • Use scalars, tf.zeros_like, or tensor-shaped replacements depending on the use case.
  • The same approach works for vectors, matrices, and batches.
  • Cleaning input data early is often better than letting NaN values propagate into the model.
  • Replacement is useful, but it is still worth understanding why the NaN values appeared in the first place.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.