def _sqrtx2p1(x):
  """Implementation of `sqrt(1 + x**2)` which is stable despite large `x`."""
  sqrt_eps = np.sqrt(np.finfo(dtype_util.as_numpy_dtype(x.dtype)).eps)
  return tf.where(
      tf.abs(x) * sqrt_eps <= 1.,
      tf.sqrt(x**2. + 1.),
      # For large x, calculating x**2 can overflow. This can be alleviated by
      # considering:
      # sqrt(1 + x**2)
      # = exp(0.5 log(1 + x**2))
      # = exp(0.5 log(x**2 * (1 + x**-2)))
      # = exp(log(x) + 0.5 * log(1 + x**-2))
      # = |x| * exp(0.5 log(1 + x**-2))
      # = |x| * sqrt(1 + x**-2)
      # We omit the last term in this approximation.
      # When |x| > 1 / sqrt(machineepsilon), the second term will be 1,
      # due to sqrt(1 + x**-2) = 1. This is also true with the gradient term,
      # and higher order gradients, since the first order derivative of
      # sqrt(1 + x**-2) is -2 * x**-3 / (1 + x**-2) = -2 / (x**3 + x),
      # and all nth-order derivatives will be O(x**-(n + 2)). This makes any
      # gradient terms that contain any derivatives of sqrt(1 + x**-2) vanish.
      tf.abs(x))