code: def get_average_batch_size(config: ml_collections.ConfigDict):
      """Computes average batch size."""
  
      if config.get("batch_size") is not None:
          return config.batch_size
  
      batch_sizes_sum = 0
      n_datasets = 0
  
      for bs in config.batch_sizes.values():
          batch_sizes_sum += bs
          n_datasets += 1
  
      average_batch_size = int(batch_sizes_sum // n_datasets)
  
      return average_batch_size
  
file_path: scenic/scenic/projects/polyvit/train_utils.py
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import datasets
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 2 for calling a global function with example usage: get_average_batch_size(config) and returns: average_batch_size, config