code: def _sqrt_no_nan(x):
      """Returns square root with a gradient at 0 being 0."""
      root = tf.math.sqrt(x)
  
      def grad(upstream):
          return tf.math.divide_no_nan(upstream, root) / 2
  
      return root, grad
  
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
  from tf_quant_finance.math import piecewise
  from tf_quant_finance.math.pde import fd_solvers
  from tf_quant_finance.models import ito_process
  from tf_quant_finance.models import utils
label_desc: Use image node_id 3 for calling a global function with example usage: _sqrt_no_nan(x) and returns: root,  grad, unknown