code: def testSqrtWithFiniteGradsBackpropsCorrectly(self):
      # Part of implementing a tf.custom_gradient is correctly handling the
      # `grad_ys` value that is propagating back from downstream ops. This test
      # checks that we got this right, in a particular case where our sqrt
      # function is squashed between a couple of other functions.
      def f(x):
          return x**2
  
      def g(x):
          return util.sqrt_with_finite_grads(x)
  
      def h(x):
          return tf.sin(x) ** 2
  
      # We only test away from zero, since we know the values don't match there.
      xs = tf.constant(np.linspace(1e-10, 10.0, 100))
      _, grad_tf_sqrt = value_and_gradient(
          lambda xs_: f(tf.sqrt(h(xs_))), xs
      )
      _, grad_safe_sqrt = value_and_gradient(
          lambda xs_: f(g(h(xs_))), xs
      )
      self.assertAllClose(
          *self.evaluate([grad_tf_sqrt, grad_safe_sqrt]), rtol=1e-10
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.gradient import value_and_gradient
  from tensorflow_probability.python.math.psd_kernels.internal import util
label_desc: Use image node_id 10 for calling the UtilTest obj's underlying member method code with example usage: obj.testSqrtWithFiniteGradsBackpropsCorrectly() and returns: unknown, util, unknown