code: def testCompileInGraphMode(self):
      if not xla.is_xla_available():
          self.skipTest("Skipping test: XLA is not available.")
  
      @tf.function
      @xla.compile_in_graph_mode
      def add(x, y):
          return x + y
  
      z = add(1.0, 2.0)
      self.assertAllClose(3.0, self.evaluate(z))
  
      @tf.function
      @xla.compile_in_graph_mode
      def add_subtract(x, y):
          return {"add": x + y, "sub": x - y}
  
      z = add_subtract(1.0, 2.0)
      self.assertAllClose({"add": 3.0, "sub": -1.0}, self.evaluate(z))
  
      class Divider(object):
          def __init__(self, x):
              self._x = x
  
          @tf.function
          @xla.compile_method_in_graph_mode
          def add_divide(self, yz):
              return self._x + yz["y"] / yz["z"]
  
      o = Divider(1.0)
      z = o.add_divide(yz={"y": 2.0, "z": 3.0})
      self.assertAllClose(1.0 + 2.0 / 3.0, self.evaluate(z))
  
      if not tf.compat.v1.executing_eagerly():
          # TF2 seems to have trouble with soft device placement (both in eager and
          # tf.function mode); and here we're specifically testing what happens when
          # XLA is not available, e.g., because we didn't compile with GPU support.
          with tf.device("/gpu:0"):
              z = add_subtract(1.0, 2.0)
          self.assertAllClose(
              {"add": 3.0, "sub": -1.0}, self.evaluate(z)
          )
  
file_path: agents/tf_agents/utils/xla_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tf_agents.utils import xla
label_desc: Use image node_id 3 for calling the XLATest obj's underlying member method code with example usage: obj.testCompileInGraphMode() and returns: unknown, dict, unknown