def _pack_images(images, rows, cols):
  """Helper utility to make a tiled field of images from numpy arrays.

  Args:
    images: Image tensor in shape [N, W, H, C].
    rows: Number of images per row in tiled image.
    cols: Number of images per column in tiled image.

  Returns:
    A tiled image of shape [W * rows, H * cols, C].
    Truncates incomplete rows.
  """
  shape = np.shape(images)
  width, height, depth = shape[-3:]
  images = np.reshape(images, (-1, width, height, depth))
  batch = np.shape(images)[0]
  rows = np.minimum(rows, batch)
  cols = np.minimum(batch // rows, cols)
  images = images[:rows * cols]
  images = np.reshape(images, (rows, cols, width, height, depth))
  images = np.transpose(images, [0, 2, 1, 3, 4])
  images = np.reshape(images, [rows * width, cols * height, depth])
  return images
Helper utility to make a tiled field of images from numpy arrays.

  Args:
    images: Image tensor in shape [N, W, H, C].
    rows: Number of images per row in tiled image.
    cols: Number of images per column in tiled image.

  Returns:
    A tiled image of shape [W * rows, H * cols, C].
    Truncates incomplete rows.
  
trax/jaxboard.py