Converting all nan values to zero in tensforflow

Question

I am trying to convert all nan values to zero in my final results. I am not able to execute it properly!

The following is the code: also availvable on colab : link

import tensorflow as tf
import numpy as np

ts = tf.constant([[0,0]]) 
tx = tf.constant([[0,1]]) 

out = ts / (ts + fx)

out.numpy() # array([nan,  0.])

tf.math.is_nan(out).numpy() # array([ True, False]

out.numpy()[(tf.math.is_nan(out).numpy())] = 0
out.numpy() #array([nan,  0.])

out.numpy() should give array([0., 0.])

Answer 1

Change:

out.numpy()[(tf.math.is_nan(out).numpy())] = 0

To:

out = tf.where(tf.math.is_nan(out), 0., out)