code: def sqrt1pm1(x):
      """Compute `sqrt(x + 1) - 1` elementwise in a numerically stable way.
  
      Args:
        x: Float `Tensor`.
  
      Returns:
        sqrt1pm1: Float `Tensor` of `sqrt1pm1(x)`.
      """
      # We follow Boost
      # https://www.boost.org/doc/libs/1_49_0/libs/math/doc/sf_and_dist/html/math_toolkit/special/powers/sqrt1pm1.html
      # and compute expm1(0.5 * log1p(x)).
      #
      # We can also derive an alternative formula by multiplying and
      # dividing by sqrt(x + 1) + 1:
      #   sqrt(x + 1) - 1 = (x + 1 - 1) / (sqrt(x + 1) + 1)
      #                   = x / (sqrt(x + 1) + 1)
      # The latter form is well-conditioned everywhere, and in particular
      # does not experience catastrophic cancellation when x ~ 0.  However,
      # without where-gating, it emits `nan` when x is `+inf`.
      return tf.math.expm1(0.5 * tf.math.log1p(x))
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import variadic_reduce
  from tensorflow_probability.python.math.scan_associative import scan_associative
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 16 for calling a global function with example usage: sqrt1pm1(x) and returns: tf