code: def _run_tf_function(func_tf: Callable, *tf_args, mode: str):
      if mode == "eager":
          return func_tf(*tf_args)  # EAGER
      elif mode == "graph":
          return tf.function(  # GRAPH
              func_tf,
              autograph=False,
              # Note that jit_compile defaults to True on TPU and False elsewhere
              input_signature=_make_tf_input_signature(*tf_args),
          )(
              *tf_args
          )  # GRAPH
      elif mode == "compiled":
          # Adding an explicit input_signature prevents TF from constant-folding
          # the computation eagerly before compilation
          return tf.function(  # COMPILED
              func_tf,
              autograph=False,
              jit_compile=True,
              input_signature=_make_tf_input_signature(*tf_args),
          )(
              *tf_args
          )  # COMPILED
      else:
          assert (
              False
          ), f"Expected 'eager', 'graph', or 'compiled' for mode: got '{mode}'"
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import dataclasses
  import re
  import os
  from typing import Any, Callable, Optional
  from absl.testing import absltest
  from absl import logging
  import jax
  from jax import dtypes
  from jax import numpy
  from jax._src import test_util
  from jax import tree_util
  from jax.experimental import jax2tf
  from jax.experimental import export
  from jax._src import config
  from jax._src import xla_bridge
  import numpy
  import tensorflow
  from tensorflow.compiler.xla import xla_data_pb2
  from tensorflow.compiler.tf2xla.python import xla
label_desc: Use image node_id 2 for calling a global function with example usage: _run_tf_function(func_tf) and returns: func_tf, tf, tf