code: def generate_image_grid_from_video(
      target: jnp.ndarray,
      prediction: jnp.ndarray,
      prediction_masks: jnp.ndarray,
      patch_size: Tuple[int, int, int],
      input_size: Tuple[int, int, int, int],
      select_central_frame: bool,
      n_columns: int = 1,
      num_img_in_column: int = 1,
  ):
      """Generates a grid of images for a video for logging summaries.
  
      When select_central_frame=True, all the reconstructed frames
      ( = video length / temporal size) will be stacked and shown onto a single row.
      When select_central_frame=False, the reconstructed frames
      ( = video length) will be split into multiple rows by num_img_in_column. It is
      highly recommended to set n_columns to 1 if the video is long.
  
      Args:
       target: Array of shape [batch, num_tokens, patch_size * patch_size *
         channels] or [batch, num_tokens, patch_size_t * patch_size_h * patch_size_w
         channels] depending on the value of select_central_frame.
         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.
       input_size: The size of the input TxHxWxC without the batch dimension.
       select_central_frame: If only the central frame is used for reconstruction.
       n_columns: The number of columns to generate in the summary. When
       select_central_frame=False is highly recommended to set it to 1.
       num_img_in_column: This is used only when select_central_frame=False. It
       splits the frames from an example onto multiple rows instead of stacking them
       onto a single row. It must be a multiple of the temporal size.
      Returns:
        An image grid. The shape is [height, 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(lgeorgescu): 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.')
      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:
          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_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_normalised + pred_mask * prediction_normalised
  
      batch_size = target.shape[0]
      if select_central_frame:
          channels = int(
              target.shape[2] / (patch_size[0] * patch_size[1])
          )
      else:
          channels = int(
              target.shape[2]
              / (patch_size[0] * patch_size[1] * patch_size[2])
          )
  
      height = input_size[1]
      width = input_size[2]
      temporal_dims = input_size[0] // patch_size[2]
      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]
  
          if select_central_frame:
              # Reshape the tensor as expected by the `patch_image` function.
              tensor_reshaped = jnp.reshape(
                  tensor,
                  (
                      batch_size,
                      temporal_dims,
                      n_patches_h,
                      n_patches_w,
                      patch_size[0],
                      patch_size[1],
                      3,
                  ),
              )
  
              images_list = jax.vmap(
                  functools.partial(
                      nn_ops.patch_image,
                      inputs_shape=(batch_size, height, width, 3),
                      patch_size=patch_size[:2],
                      mode="p2i",
                  ),
                  in_axes=1,
                  out_axes=0,
              )(tensor_reshaped)
  
              final_image = jnp.concatenate(images_list, axis=2)
              return final_image
  
          else:
              # Reshape the tensor as expected by the `patch_image` function.
              tensor_reshaped = jnp.reshape(
                  tensor,
                  (
                      batch_size,
                      temporal_dims,
                      n_patches_h,
                      n_patches_w,
                      patch_size[2],
                      patch_size[0],
                      patch_size[1],
                      3,
                  ),
              )
              images_list = []
              for temporal_video_idx in range(temporal_dims):
                  images_list_patch = jax.vmap(
                      functools.partial(
                          nn_ops.patch_image,
                          inputs_shape=(batch_size, height, width, 3),
                          patch_size=patch_size[:2],
                          mode="p2i",
                      ),
                      in_axes=3,
                      out_axes=0,
                  )(tensor_reshaped[:, temporal_video_idx])
  
                  images_list.extend(images_list_patch)
  
              if num_img_in_column % patch_size[2] != 0:
                  raise ValueError(
                      "`patch_size[2]` must divide"
                      "the `num_img_in_column` evenly!"
                  )
  
              grouped_images_list = np.array_split(
                  images_list,
                  max(1, len(images_list) // num_img_in_column),
              )
  
              final_image = jnp.concatenate(
                  [
                      jnp.concatenate(sub_images_list, axis=2)
                      for sub_images_list in grouped_images_list
                  ],
                  axis=1,
              )
              return final_image
  
      input_images = unpatch_to_images(input_normalised)
      pred_images = unpatch_to_images(prediction_normalised)
      masked_token_images = unpatch_to_images(pred_masked_tokens)
      unmasked_token_images = unpatch_to_images(pred_unmasked_tokens)
      target_images = unpatch_to_images(target_normalised)
      pred_with_unmasked_replaced_images = unpatch_to_images(
          pred_with_unmasked_replaced
      )
  
      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(
                  target_images,
                  input_images,
                  pred_images,
                  masked_token_images,
                  unmasked_token_images,
                  pred_with_unmasked_replaced_images,
              )
          ]
      )
      # add margin 2% of the image height
      margin_size = max(int(0.02 * images_concat.shape[1]), 10)
      margin = (
          np.ones(
              (
                  images_concat.shape[0],
                  margin_size,
                  images_concat.shape[2],
                  images_concat.shape[3],
              ),
              dtype=np.uint8,
          )
          * 255
      )
      images_concat = jax.numpy.concatenate(
          (images_concat, margin), axis=1
      )
  
      new_height = unmasked_token_images.shape[1]
      new_width = unmasked_token_images.shape[2]
      image_grid = images_concat.reshape(
          n_rows,
          n_columns,
          6 * new_height + margin_size,
          new_width,
          channels,
      ).swapaxes(1, 2)
      image_grid = image_grid.reshape(
          (6 * new_height * n_rows + batch_size * margin_size),
          new_width * n_columns,
          channels,
      ).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 2 for calling a global function with example usage: generate_image_grid_from_video(target, prediction, prediction_masks, patch_size, input_size, select_central_frame, n_columns, num_img_in_column) and returns: image_grid, unknown, final_image, final_image