code: def generate_image_grid(
      target: jnp.ndarray,
      prediction: jnp.ndarray,
      prediction_masks: jnp.ndarray,
      patch_size: Tuple[int, int],
      n_columns: int,
      input_size: Tuple[int, int, int],
      modality: str = "rgb",
  ):
      """Generates a grid of images for logging summaries.
  
      Args:
       target: Array of shape [batch, num_tokens, patch_size * patch_size *
         channels]. Values are in the range [-1, 1].
       prediction: Array of same shape as target. It is possible that values are
         outside of the range [-1, 1].
       prediction_masks: Array of shape [batch, num_tokens]. Each entry is binary,
         ie 0 or 1. Used to mask out patches in the input and predictions.
       patch_size: The patch size of the model.
       n_columns: The number of columns to generate in the summary.
       input_size: The size of the input, specified as [height, width, channels].
       modality: The modality which is shown.
      Returns:
        An image grid. The shape is
        [batch_size / n_columns * 2 * height, n_columns * width, channels].
        Data type is uint8, and values are in the range [0, 255].
      """
      if target.shape != prediction.shape:
          raise ValueError(
              "Shape of target and prediction must match."
              f"{target.shape} vs {prediction.shape}."
          )
      # TODO(aarnab): Correctly handle the case where the input is normalised by
      # mean and std-dev instead of [-1, 1].
      # if jnp.max(target) > 1 or jnp.min(target) < -1:
      #   raise ValueError('Invalid ranges in target.')
      if modality == "spectrogram":
          prediction_clipped = prediction
      else:
          prediction_clipped = jnp.clip(prediction, a_min=-1, a_max=1)
  
      # Normalise to uint8 in range [0, 255] for summary-writing.
      def normalise(
          tensor: jnp.ndarray, offset: float = 127.5
      ) -> jnp.ndarray:
          if modality == "spectrogram":
              return tensor
          else:
              return tensor * offset + offset
  
      target_normalised = normalise(target)
      prediction_normalised = normalise(prediction_clipped)
  
      # if prediction_masks is not None:
      # Mask out corresponding part of the input.
      mask_image = jnp.zeros(target.shape)
      pred_mask = jnp.expand_dims(prediction_masks, axis=-1)
      input_normalized = target_normalised
  
      target_normalised = (
          1 - pred_mask
      ) * target_normalised + pred_mask * mask_image
      pred_unmasked_tokens = (
          1 - pred_mask
      ) * prediction_normalised + pred_mask * mask_image
      pred_masked_tokens = (
          1 - pred_mask
      ) * mask_image + pred_mask * prediction_normalised
      pred_with_unmasked_replaced = (
          1 - pred_mask
      ) * input_normalized + pred_mask * prediction_normalised
  
      batch_size = target.shape[0]
      height, width, channels = input_size
      n_rows = batch_size // n_columns
      if batch_size != n_rows * n_columns:
          raise ValueError(
              "`n_columns` must divide the number of images evenly"
          )
  
      def unpatch_to_images(tensor: jnp.ndarray) -> jnp.ndarray:
          n_patches_h = height // patch_size[0]
          n_patches_w = width // patch_size[1]
          # Reshape the tensor as expected by the `patch_image` function.
          tensor_reshaped = jnp.reshape(
              tensor,
              (
                  batch_size,
                  n_patches_h,
                  n_patches_w,
                  patch_size[0],
                  patch_size[1],
                  channels,
              ),
          )
          return nn_ops.patch_image(
              tensor_reshaped,
              inputs_shape=(batch_size, height, width, channels),
              patch_size=patch_size,
              mode="p2i",
          )
  
      input_images = unpatch_to_images(input_normalized)
      target_images = unpatch_to_images(target_normalised)
      prediction_images = unpatch_to_images(prediction_normalised)
      masked_token_images = unpatch_to_images(pred_masked_tokens)
      unmasked_token_images = unpatch_to_images(pred_unmasked_tokens)
      pred_with_unmasked_replaced_images = unpatch_to_images(
          pred_with_unmasked_replaced
      )
  
      num_images = 6
      if modality == "spectrogram":
          images_concat = jnp.array(
              [
                  jnp.concatenate(
                      [x1, x2, x3, x4, x5, x6], axis=1
                  )  # pylint: disable= g-complex-comprehension
                  for x1, x2, x3, x4, x5, x6 in zip(
                      input_images,
                      target_images,
                      prediction_images,
                      masked_token_images,
                      unmasked_token_images,
                      pred_with_unmasked_replaced_images,
                  )
              ]
          )
          image_grid = images_concat.reshape(
              n_rows, n_columns, height, num_images * width, channels
          ).swapaxes(1, 2)
  
          image_grid = image_grid.reshape(
              height * n_rows, num_images * width * n_columns, channels
          )
      else:
          images_concat = jnp.array(
              [
                  jnp.concatenate(
                      [x1, x2, x3, x4, x5, x6], axis=0
                  )  # pylint: disable= g-complex-comprehension
                  for x1, x2, x3, x4, x5, x6 in zip(
                      input_images,
                      target_images,
                      prediction_images,
                      masked_token_images,
                      unmasked_token_images,
                      pred_with_unmasked_replaced_images,
                  )
              ]
          )
  
          image_grid = images_concat.reshape(
              n_rows, n_columns, num_images * height, width, channels
          ).swapaxes(1, 2)
  
          image_grid = image_grid.reshape(
              num_images * height * n_rows, width * n_columns, channels
          ).astype(jnp.uint8)
  
      if modality == "spectrogram":
          # Normalize the entire image
          image_grid = image_grid - image_grid.min()
          image_grid = image_grid / image_grid.max()
  
          cm = plt.get_cmap("viridis")
          # plt.get_cmap expects uint8 as input
          image_grid = image_grid * 255
          image_grid = cm(image_grid[:, :, 0].astype(jnp.uint8))
          # image_grid is between [0, 1]
          image_grid = (image_grid[:, :, :3] * 255).astype(jnp.uint8)
  
      return image_grid
  
imports: import functools
  from typing import Any, Callable, Dict, Optional, Tuple, Union
  from absl import logging
  from clu import platform
  import flax.linen
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.layers import nn_ops
  from scenic.train_lib import train_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: generate_image_grid(target, prediction, prediction_masks, patch_size, n_columns, input_size, modality) and returns: image_grid, nn_ops, tensor, unknown