code: def tf_accelerator_and_tolerances():
      """Picks the TF accelerator to use and the tolerances for numerical checks."""
      tf_accelerator = (
          tf.config.list_logical_devices("TPU")
          + tf.config.list_logical_devices("GPU")
          + tf.config.list_logical_devices("CPU")
      )[0]
      logging.info("Using tf_accelerator = %s", tf_accelerator)
      if tf_accelerator.device_type == "TPU":
          tolerances = dict(atol=1e-6, rtol=1e-6)
      elif tf_accelerator.device_type == "GPU":
          tolerances = dict(atol=1e-6, rtol=1e-4)
      elif tf_accelerator.device_type == "CPU":
          tolerances = dict(atol=1e-5, rtol=1e-5)
      logging.info("Using tolerances %s", tolerances)
      return tf_accelerator, tolerances
  
file_path: jax/jax/experimental/jax2tf/examples/saved_model_main.py
imports: import logging
  import os
  from absl import app
  from absl import flags
  from jax.experimental.jax2tf.examples import mnist_lib
  from jax.experimental.jax2tf.examples import saved_model_lib
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 5 for calling a global function with example usage: tf_accelerator_and_tolerances() and returns: tf_accelerator,  tolerances