code: def sqrt_with_finite_grads(x, name=None):
      """A sqrt function whose gradient at zero is very large but finite.
  
      Args:
        x: a `Tensor` whose sqrt is to be computed.
        name: a Python `str` prefixed to all ops created by this function.
          Default `None` (i.e., "sqrt_with_finite_grads").
  
      Returns:
        sqrt: the square root of `x`, with an overridden gradient at zero
        grad: a gradient function, which is the same as sqrt's gradient everywhere
          except at zero, where it is given a large finite value, instead of `inf`.
  
      Raises:
        TypeError: if `tf.convert_to_tensor(x)` is not a `float` type.
  
      Often in kernel functions, we need to compute the L2 norm of the difference
      between two vectors, `x` and `y`: `sqrt(sum_i((x_i - y_i) ** 2))`. In the
      case where `x` and `y` are identical, e.g., on the diagonal of a kernel
      matrix, we get `NaN`s when we take gradients with respect to the inputs. To
      see, this consider the forward pass:
  
        ```
        [x_1 ... x_N]  -->  [x_1 ** 2 ... x_N ** 2]  -->
            (x_1 ** 2 + ... + x_N ** 2)  -->  sqrt((x_1 ** 2 + ... + x_N ** 2))
        ```
  
      When we backprop through this forward pass, the `sqrt` yields an `inf` because
      `grad_z(sqrt(z)) = 1 / (2 * sqrt(z))`. Continuing the backprop to the left, at
      the `x ** 2` term, we pick up a `2 * x`, and when `x` is zero, we get
      `0 * inf`, which is `NaN`.
  
      We'd like to avoid these `NaN`s, since they infect the rest of the connected
      computation graph. Practically, when two inputs to a kernel function are
      equal, we are in one of two scenarios:
        1. We are actually computing k(x, x), in which case norm(x - x) is
           identically zero, independent of x. In this case, we'd like the
           gradient to reflect this independence: it should be zero.
        2. We are computing k(x, y), and x just *happens* to have the same value
           as y. The gradient at such inputs is in fact ill-defined (there is a
           cusp in the sqrt((x - y) ** 2) surface along the line x = y). There are,
           however, an infinite number of sub-gradients, all of which are valid at
           all such inputs. By symmetry, there is exactly one which is "special":
           zero, and we elect to use that value here. In practice, having two
           identical inputs to a kernel matrix is probably a pathological
           situation to be avoided, but that is better resolved at a higher level
           than this.
  
      To avoid the infinite gradient at zero, we use tf.custom_gradient to redefine
      the gradient at zero. We assign it to be a very large value, specifically
      the sqrt of the max value of the floating point dtype of the input. We use
      the sqrt (as opposed to just using the max floating point value) to avoid
      potential overflow when combining this value with others downstream.
      """
      with tf.name_scope(name or "sqrt_with_finite_grads"):
          x = tf.convert_to_tensor(value=x, name="x")
          if not dtype_util.is_floating(x.dtype):
              raise TypeError("Input `x` must be floating type.")
  
          def grad(grad_ys):
              large_float_like_x = np.sqrt(
                  np.finfo(dtype_util.as_numpy_dtype(x.dtype)).max
              )
              safe_grads = tf.where(
                  tf.equal(x, 0),
                  large_float_like_x,
                  0.5 * tf.math.rsqrt(x),
              )
              return grad_ys * safe_grads
  
          return tf.sqrt(x), grad
  
file_path: probability/tensorflow_probability/python/math/psd_kernels/internal/util.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 3 for calling a global function with example usage: sqrt_with_finite_grads(x, name) and returns: tf,  grad, unknown