code: def get_batch_size(stage_id, **kwargs):
      """Returns batch size for each stage.
  
      It is expected that `len(batch_size_schedule) == num_resolutions`. Each stage
      corresponds to a resolution and hence a batch size. However if
      `len(batch_size_schedule) < num_resolutions`, pad `batch_size_schedule` in the
      beginning with the first batch size.
  
      Args:
        stage_id: An integer of training stage index.
        **kwargs: A dictionary of
            'batch_size_schedule': A list of integer, each element is the batch size
                for the current training image resolution.
            'num_resolutions': An integer of number of progressive resolutions.
  
      Returns:
        An integer batch size for the `stage_id`.
      """
      batch_size_schedule = kwargs["batch_size_schedule"]
      num_resolutions = kwargs["num_resolutions"]
      if len(batch_size_schedule) < num_resolutions:
          batch_size_schedule = [batch_size_schedule[0]] * (
              num_resolutions - len(batch_size_schedule)
          ) + batch_size_schedule
  
      return int(batch_size_schedule[(stage_id + 1) // 2])
  
imports: import os
  import time
  from absl import logging
  from magenta.models.gansynth.lib import networks
  import numpy
  import tensorflow.compat.v1
  import tensorflow_gan
label_desc: Use image node_id 5 for calling a global function with example usage: get_batch_size(stage_id) and returns: int