def draw_frames_side_by_side(*video_batches: jnp.ndarray,
                             show_num: int = 4,
                             border: int = 1):
  """Draws the original and generated frames side-by-side."""
  assert all([v.ndim == 5 for v in video_batches]), 'Incorrect input format.'
  swap_ins = 'device_bs t h w c -> device_bs h (t w) c'
  vs = [einops.rearrange(v, swap_ins) for v in video_batches]
  v = np.concatenate(vs, 1)
  if border > 0:
    v = np.pad(v, ((0, 0), (border, border), (border, border), (0, 0)))
  return make_grid(v, show_num)
Draws the original and generated frames side-by-side.
videogvt/train_lib/train_utils.py