code: def main(argv):
      if len(argv) > 1:
          raise app.UsageError("Too many command-line arguments.")
  
      # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make
      # it unavailable to JAX.
      tf.config.experimental.set_visible_devices([], "GPU")
  
      logging.info(
          "JAX process: %d / %d",
          jax.process_index(),
          jax.process_count(),
      )
      logging.info("JAX local devices: %r", jax.local_devices())
  
      # Add a note so that we can tell which task is which JAX host.
      # (Depending on the platform task 0 is not guaranteed to be host 0)
      platform.work_unit().set_task_status(
          f"process_index: {jax.process_index()}, "
          f"process_count: {jax.process_count()}"
      )
      platform.work_unit().create_artifact(
          platform.ArtifactType.DIRECTORY, FLAGS.workdir, "workdir"
      )
  
      train.train_and_evaluate(FLAGS.config, FLAGS.workdir)
  
imports: import jax
  import tensorflow
  import train
  from absl import app, flags, logging
  from clu import platform
  from ml_collections import config_flags
label_desc: Use image node_id 1 for calling a global function with example usage: main(argv) without return types