code: def main(_):
      assert FLAGS.model_dirs_or_checkpoints
  
      if not tf.io.gfile.exists(FLAGS.output_dir):
          tf.io.gfile.makedirs(FLAGS.output_dir)
  
      if (
          FLAGS.operation == "average_last_n"
          and len(FLAGS.model_dirs_or_checkpoints) > 1
      ):
          raise ValueError(
              "Need only 1 directory for %s operation" % FLAGS.operation
          )
  
      checkpoints = []
  
      for path in FLAGS.model_dirs_or_checkpoints:
          if tf.io.gfile.isdir(path):
              # Grab the latest checkpoint for all the provided model dirs
              checkpoint_state = tf.train.get_checkpoint_state(path)
              if FLAGS.operation == "average_last_n":
                  ckpt_paths = tf.io.gfile.glob(
                      os.path.join(path, "model.ckpt*index")
                  )
  
                  def sort_fn(ckpt):
                      return int(re.sub(".*ckpt-", "", ckpt))
  
                  ckpts = sorted(
                      [c.replace(".index", "") for c in ckpt_paths],
                      key=sort_fn,
                  )
                  checkpoints.extend(
                      ckpts[-FLAGS.number_of_checkpoints :]
                  )
              else:
                  checkpoints.append(
                      checkpoint_state.all_model_checkpoint_paths[-1]
                  )
          else:
              if FLAGS.operation == "average_last_n":
                  raise ValueError(
                      "need a directory while running %s operation"
                      % FLAGS.operation
                  )
              checkpoints.append(path)
  
      logging.info("Using checkpoints %s", checkpoints)
  
      if FLAGS.operation in ["ensemble", "average", "average_last_n"]:
          if len(checkpoints) == 1:
              raise ValueError(
                  "no point in ensebling/averaging one checkpoint"
              )
      else:
          if len(checkpoints) != 1:
              raise ValueError(
                  "operation %s requires exactly one checkpoint"
                  % FLAGS.operation
              )
  
      var_values = {}
      var_dtypes = {}
  
      for i in range(0, len(checkpoints)):
          checkpoint = checkpoints[i]
          logging.info("loading checkpoint %s", checkpoint)
          reader = tf.train.load_checkpoint(checkpoint)
          var_list = tf.train.list_variables(checkpoint)
          for name, _ in var_list:
              if i:
                  assert name in var_values
                  tensor = reader.get_tensor(name)
                  assert tensor.dtype == var_dtypes[name]
                  var_values[name].append(tensor)
              else:
                  tensor = reader.get_tensor(name)
                  var_dtypes[name] = tensor.dtype
                  var_values[name] = [tensor]
                  if not FLAGS.global_step:
                      if name == "global_step":
                          FLAGS.global_step = tensor
  
          logging.info("Read from checkpoint %s", checkpoint)
  
      new_var_values = {}
  
      # stack the list of tensors along the 0th dimension.
      for name, tensors in var_values.items():
          tensor = tensors[0]
          if name == "global_step":
              new_val = np.int32(FLAGS.global_step)
          elif FLAGS.operation == "ensemble":
              new_val = np.stack(tensors)
          elif FLAGS.operation == "autoensemble":
              new_val = np.stack([tensor] * FLAGS.autoensemble_size)
          elif (
              FLAGS.operation == "average"
              or FLAGS.operation == "average_last_n"
          ):
              new_val = average_tensors(tensors)
          elif FLAGS.operation == "extract_first":
              new_val = tensor[0]
          else:
              raise ValueError(
                  "unknown FLAGS.operation=%s" % FLAGS.operation
              )
          new_var_values[name] = new_val
  
      var_values = new_var_values
  
      with tf.variable_scope(
          tf.get_variable_scope(), reuse=tf.AUTO_REUSE
      ):
          tf_vars = [
              tf.get_variable(
                  v, shape=var_values[v].shape, dtype=var_dtypes[v]
              )
              for v in var_values
          ]
  
      placeholders = [
          tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars
      ]
      assign_ops = [
          tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)
      ]
      saver = tf.train.Saver(tf.all_variables())
  
      output_file = "model.ckpt-" + str(FLAGS.global_step)
      output_path = os.path.join(FLAGS.output_dir, output_file)
  
      # Build a model consisting only of variables, set them to the average values.
      with tf.Session() as sess:
          sess.run(tf.global_variables_initializer())
          for p, assign_op, (name, value) in zip(
              placeholders, assign_ops, var_values.items()
          ):
              sess.run(assign_op, {p: value})
          # Use the built saver to save the averaged checkpoint.
          saver.save(sess, output_path)
  
      logging.info("Transformed checkpoints saved in %s", output_path)
  
file_path: text-to-text-transfer-transformer/t5/scripts/transform_checkpoints.py
imports: import os
  import re
  from absl import app
  from absl import flags
  from absl import logging
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: main(_) and returns: int