def _is_finite(x: common_types.TensorType) -> common_types.TensorType:
  """Extension of `tf.math.is_finite` that works with all dtypes."""
  if x.dtype.is_floating:
    return tf.math.is_finite(x)
  return tf.ones_like(x, dtype=tf.bool)
Extension of `tf.math.is_finite` that works with all dtypes.
tensorflow_transform/tf_utils.py