code: def testDerivative(self):
      def fn(x, a):
          return 2 * x + a
  
      with ops.Graph().as_default() as g:
          xla_func = def_function.function(fn, jit_compile=True)
          with backprop.GradientTape() as tape:
              inputs = array_ops.placeholder(dtypes.float32, [5])
              tape.watch(inputs)
              outputs = xla_func(inputs, 1)
          grads = tape.gradient(outputs, inputs)
  
      with session.Session(graph=g) as sess:
          grads_tensor = sess.run(
              grads, feed_dict={inputs: [1, 2, 2, 3, 3]}
          )
          self.assertAllClose([2, 2, 2, 2, 2], grads_tensor)
          (forward, backward) = xla_func.get_concrete_function(
              inputs, 1
          )._delayed_rewrite_functions.forward_backward()
  
          # Check that the must-compile attribute gets correctly propagated to the
          # created derivatives.
          self.assertTrue(
              forward.cached_definition.attr["_XlaMustCompile"]
          )
          self.assertTrue(backward.function_def.attr["_XlaMustCompile"])
  
file_path: tensorflow/tensorflow/python/compiler/xla/tests/jit_compile_test.py
imports: from tensorflow.python.client import session
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import string_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling the JitCompileTest obj's underlying member method code with example usage: obj.testDerivative() and returns: unknown