docstr: Test for the custom square root.
docstr_tok: ["Test","for","the","custom","square","root","."]
code: def test_sqrt_grad(self):
      """Test for the custom square root."""
      dtype = tf.float64
      x1 = tf.random.stateless_uniform(
          shape=[10, 5], seed=[1, 2], dtype=dtype
      )
      x2 = tf.constant([0.0, 1.0], dtype=dtype)
      with tf.GradientTape(persistent=True) as tape:
          tape.watch([x1, x2])
          y1 = univariate_geometric_brownian_motion._sqrt_no_nan(x1)
          y2 = univariate_geometric_brownian_motion._sqrt_no_nan(x2)
          y1_true = tf.sqrt(x1)
  
      with self.subTest("Value"):
          self.assertAllClose(y1, y1_true)
  
      with self.subTest("GradientCorrect"):
          self.assertAllClose(
              tape.gradient(y1, x1), tape.gradient(y1_true, x1)
          )
  
      with self.subTest("ZeroGradientCorrect"):
          self.assertAllClose(tape.gradient(y2, x2), [0.0, 0.5])
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
  from tf_quant_finance.models.geometric_brownian_motion import geometric_brownian_motion_test_utils
  from tf_quant_finance.models.geometric_brownian_motion import univariate_geometric_brownian_motion
label_desc: Use image node_id 26 for calling the GeometricBrownianMotionTest obj's underlying member method code with example usage: obj.test_sqrt_grad() without return types