code: def compile_in_graph_mode(fn):
      """Decorator for XLA compilation iff in graph mode and XLA is available.
  
      Example:
  
      ```python
      @compile_in_graph_mode
      def fn(x, y, z):
        return {'a': x + y, 'b': y * z}
  
      @common.function
      def calls_fn(inputs):
        return fn(inputs.x, inputs.y, inputs.z)
  
      # Call calls_fn().
  
      Args:
        fn: A callable that accepts a list of possibly nested tensor arguments.
          kwargs and inputs taking the value `None` are not supported.  Non-tensor
          arguments are treated as nest objects, and leaves are converted to
          tensors.
  
      Returns:
        A function that, when called, checks if XLA is compiled in and enabled
        for the current device, and that it's being built in graph mode, and
        returns an XLA-compiled version of `fn`.  If in eager mode, or XLA
        is not available, then `fn` is called directly.
      ```
      """
  
      @functools.wraps(fn)
      def _call_compiled(*args, **kwargs):
          return _compiled(*args, _fn=fn, _self=None, **kwargs)
  
      return _call_compiled
  
file_path: agents/tf_agents/utils/xla.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: compile_in_graph_mode(fn) and returns: _call_compiled, _compiled