code: def test_correspondence(
      self, x, new_fn, old_fn, value_rtol=1e-4, grad_rtol=1e-4
  ):
      if FLAGS.compute_on == "tpu":
          self.skipTest(
              "The manual functions are more precise than TPU ones."
          )
      x = x.astype(dtype_util.as_numpy_dtype(self.dtype))
  
      y1 = self.run_fn(old_fn, (x,))
      y2 = self.run_fn(new_fn, (x,))
  
      self.assertAllClose(y1, y2, rtol=value_rtol)
  
      if NUMPY_MODE:
          return
  
      num_outputs = np.prod(y1.shape).astype(np.int32)
  
      dy = tf.reshape(
          tf.range(1, 1 + num_outputs, dtype=y1.dtype), y1.shape
      )
  
      y1, grads1 = self.run_fn(
          lambda x, dy: gradient.value_and_gradient(  # pylint: disable=g-long-lambda
              old_fn, x, output_gradients=dy
          ),
          (x, dy),
      )
      y2, grads2 = self.run_fn(
          lambda x, dy: gradient.value_and_gradient(  # pylint: disable=g-long-lambda
              new_fn, x, output_gradients=dy
          ),
          (x, dy),
      )
  
      self.assertAllClose(y1, y2, rtol=value_rtol)
      self.assertAllClose(grads1, grads2, rtol=grad_rtol)
  
file_path: probability/tensorflow_probability/python/experimental/math/manual_special_functions_test.py
imports: import functools
  from absl import flags
  from absl.testing import parameterized
  import mock
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.math import manual_special_functions
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 3 for calling the ManualSpecialFunctionsTest obj's underlying member method code with example usage: obj.test_correspondence(x, new_fn, old_fn, value_rtol, grad_rtol) without return types