code: def get_image_grid(
      targets: jnp.ndarray,
      predictions: jnp.ndarray,
      token_mask: jnp.ndarray,
      config: ml_collections.ConfigDict,
      input_size: Optional[
          Union[Tuple[int, int, int], Tuple[int, int, int, int]]
      ] = None,
  ) -> Optional[jnp.ndarray]:
      """Returns an image grid for summary writing."""
  
      image_grid = None
      feature_target = config.masked_feature_loss.target
  
      n_columns = config.masked_feature_loss.get(
          "summary_num_columns", 1
      )
  
      if feature_target == avmae_base_model.FeatureTargets.RGB:
          if len(config.model.patches.size) == 2:
              image_grid = avmae_train_utils.generate_image_grid(
                  targets,
                  predictions,
                  token_mask,
                  tuple(config.model.patches.size),
                  n_columns=n_columns,
                  input_size=input_size,
              )
          elif len(config.model.patches.size) == 3:
              num_img_in_column = config.masked_feature_loss.get(
                  "number_of_img_in_column", 16
              )
              select_central_frame = config.masked_feature_loss.get(
                  "select_central_frame", True
              )
  
              assert (
                  input_size is not None
              ), "Input size must be provided for video!"
  
              image_grid = (
                  avmae_train_utils.generate_image_grid_from_video(
                      targets,
                      predictions,
                      token_mask,
                      tuple(config.model.patches.size),
                      input_size,
                      n_columns=n_columns,
                      num_img_in_column=num_img_in_column,
                      select_central_frame=select_central_frame,
                  )
              )
  
      return image_grid
  
imports: import functools
  from typing import Any, Callable, Dict, Iterator, Optional, Tuple, Type, Union
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.av_mae import base_model
  from scenic.projects.av_mae import train_utils
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 5 for calling a global function with example usage: get_image_grid(targets, predictions, token_mask, config, input_size) and returns: image_grid