code: def define_device(tpu=True):
      """Register device specific flags.
      Args:
        tpu: Create flags to specify TPU operation.
      Returns:
        A list of flags for core.py to marks as key flags.
      """
  
      key_flags = []
  
      if tpu:
          flags.DEFINE_string(
              name="tpu",
              default=None,
              help=help_wrap(
                  "The Cloud TPU to use for training. This should be either the name "
                  "used when creating the Cloud TPU, or a "
                  "grpc://ip.address.of.tpu:8470 url. Passing `local` will use the"
                  "CPU of the local instance instead. (Good for debugging.)"
              ),
          )
          key_flags.append("tpu")
  
          flags.DEFINE_string(
              name="tpu_zone",
              default=None,
              help=help_wrap(
                  "[Optional] GCE zone where the Cloud TPU is located in. If not "
                  "specified, we will attempt to automatically detect the GCE "
                  "project from metadata."
              ),
          )
  
          flags.DEFINE_string(
              name="tpu_gcp_project",
              default=None,
              help=help_wrap(
                  "[Optional] Project name for the Cloud TPU-enabled project. If not "
                  "specified, we will attempt to automatically detect the GCE "
                  "project from metadata."
              ),
          )
  
          flags.DEFINE_integer(
              name="num_tpu_shards",
              default=8,
              help=help_wrap("Number of shards (TPU chips)."),
          )
  
      return key_flags
  
file_path: autodist/examples/benchmark/utils/flags/_device.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl import flags
  import tensorflow
  from utils.flags._conventions import help_wrap
label_desc: Use image node_id 2 for calling a global function with example usage: define_device(tpu) and returns: key_flags