code: def _jax_and_tf_configure_for_devices():  # pylint: disable=missing-function-docstring
      if FLAGS.use_tpu:
          jax.config.update("jax_platform_name", "tpu")
          jax.config.update("jax_xla_backend", FLAGS.jax_xla_backend)
          jax.config.update(
              "jax_backend_target", FLAGS.jax_backend_target
          )
      if FLAGS.enable_eager_execution and (
          fastmath.is_backend(Backend.NUMPY)
          or fastmath.is_backend(Backend.JAX)
      ):
          # Numpy backend doesn't benefit from having the input pipeline run on GPU,
          # and jax backend has GPU memory contention if TF uses the GPU. Gin must be
          # set up first before determining the backend.
          tf.config.experimental.set_visible_devices([], "GPU")
  
imports: import atexit
  import datetime
  import functools
  import os
  from absl import app
  from absl import flags
  from absl import logging
  import gin
  import jax
  from jax.lib import xla_extension
  import tensorflow.compat.v2
  from trax import fastmath
  from trax import trainer_flags
  from trax.supervised import trainer_lib
  from trax.tf_numpy import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: _jax_and_tf_configure_for_devices() without return types