code: def testJVP(self):
      if not JAX_MODE:
          self.skipTest("Custom JVPs are JAX-only.")
  
      def f_vjp_fwd(x, y):
          # When a JVP is specified, this function is ignored.
          raise NotImplementedError()
  
      def f_vjp_bwd(x_y, dz):
          # When a JVP is specified, this function is ignored.
          raise NotImplementedError()
  
      def f_jvp(x_y, dx_dy):
          x, y = x_y
          dx, dy = dx_dy
          return f(x, y), 7.0 * (dx * x + dy * y)
  
      @custom_gradient.custom_gradient(
          vjp_fwd=f_vjp_fwd,
          vjp_bwd=f_vjp_bwd,
          jvp_fn=f_jvp,
      )
      def f(x, y):
          return x**2 + y**2
  
      x = tf.constant(2.0)
      y = tf.constant(3.0)
      dz = tf.constant(5.0)
  
      z1 = f(x, y)
      z2, (dx, dy) = tfp_gradient.value_and_gradient(
          f, (x, y), output_gradients=dz
      )
  
      self.assertAllClose(x**2 + y**2, z1)
      self.assertAllClose(x**2 + y**2, z2)
      self.assertAllClose(7.0 * dz * x, dx)
      self.assertAllClose(7.0 * dz * y, dy)
  
      import jax  # pylint: disable=g-import-not-at-top
  
      z3, dz2 = jax.jvp(f, (x, y), (dx, dy))
      self.assertAllClose(x**2 + y**2, z3)
      self.assertAllClose(7.0 * (dx * x + dy * y), dz2)
  
file_path: probability/tensorflow_probability/python/internal/custom_gradient_test.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  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 CustomGradientTest obj's underlying member method code with example usage: obj.testJVP() and returns: f,  unknown, unknown