code: def load_and_aggregate(
      ckpt_dir, aggregating_epochs, aggregation_method, base_model_fn
  ):
      """Returns an aggregated model from checkpoints at given epochs.
  
      Args:
        ckpt_dir: Checkpoint directory to load the models.
        aggregating_epochs: A list of epoch numbers which corresponding checkpoints
          are loaded and aggregated.
        aggregation_method: How to aggregate checkpoints. Set to "avg_weight" for
          averaging model weights, or "ensemble" for ensemble (averaging
          probablistic predicitons).
        base_model_fn: A callable which builds the model structure.
  
      Returns:
        A callable which can make inference on input.
      """
      if aggregation_method not in ("avg_weight", "ensemble"):
          raise ValueError(
              f"Unknown aggregation method: {aggregation_method}"
          )
      models = []
      for epoch in aggregating_epochs:
          logging.info("Loading model at epoch %d", epoch)
          path = os.path.join(ckpt_dir, f"ckpt-{epoch}")
          models.append(helper.load_checkpoint(path, base_model_fn))
  
      logging.info("Aggregating models using %s", aggregation_method)
      if aggregation_method == "avg_weight":
          model_weights = [mdl.get_weights() for mdl in models]
          agg_weights = [
              sum(ws) / len(models) for ws in zip(*model_weights)
          ]
          agg_model = base_model_fn()
          agg_model.set_weights(agg_weights)
          return agg_model
      else:  # ensemble
  
          @tf.function
          def ensemble_model(x):
              predictions = [
                  tf.nn.softmax(model(x), axis=-1) for model in models
              ]
              avg_pred = tf.reduce_mean(
                  tf.stack(predictions, axis=0), axis=0
              )
              return tf.math.log(avg_pred)  # return in logit space
  
          return ensemble_model
  
file_path: neural-structured-learning/research/multi_representation_adversary/multi_representation_adversary/evaluator.py
imports: import functools
  import hashlib
  import os
  import time
  from absl import logging
  import gin
  from multi_representation_adversary import data
  from multi_representation_adversary import helper
  from multi_representation_adversary import resnet
  import tensorflow.compat.v2
label_desc: Use image node_id 3 for calling a global function with example usage: load_and_aggregate(ckpt_dir, aggregating_epochs, aggregation_method, base_model_fn) and returns: agg_model, ensemble_model, tf