tf.debugging.assert_all_finite(
    x, message, name=None
)

@tf.function
def f(x):
  x = tf.debugging.assert_all_finite(x, 'Input x must be all finite')
  return x + 1

f(tf.constant([np.inf, 1, 2]))
Traceback (most recent call last):

InvalidArgumentError: ...