code: def image_grid(image: np.ndarray) -> np.ndarray:
      """Rearrange array of images (nh, hw, c, h, w) into image grid in a single image (c, nh * h, nh * w)."""
      s = image.shape
      return image.transpose([2, 0, 3, 1, 4]).reshape(
          [s[2], s[3] * s[0], s[4] * s[1]]
      )
  
file_path: objax/objax/util/image.py
imports: import io
  from typing import Union, BinaryIO, IO
  import jax.numpy
  import numpy
  from PIL import Image
  from objax.typing import JaxArray
label_desc: Use image node_id 2 for calling a global function with example usage: image_grid(image) and returns: image