code: def save_image(
      ndarray, fp, nrow=8, padding=2, pad_value=0.0, format_img=None
  ):
      """Make a grid of images and Save it into an image file.
  
      Args:
        ndarray (array_like): 4D mini-batch images of shape (B x H x W x C)
        fp:  A filename(string) or file object
        nrow (int, optional): Number of images displayed in each row of the grid.
          The final grid size is ``(B / nrow, nrow)``. Default: ``8``.
        padding (int, optional): amount of padding. Default: ``2``.
        pad_value (float, optional): Value for the padded pixels. Default: ``0``.
        format_img(Optional):  If omitted, the format to use is determined from the
          filename extension. If a file object was used instead of a filename,
          this parameter should always be used.
      """
  
      if not (
          isinstance(ndarray, jnp.ndarray)
          or (
              isinstance(ndarray, list)
              and all(isinstance(t, jnp.ndarray) for t in ndarray)
          )
      ):
          raise TypeError(
              f"array_like of tensors expected, got {type(ndarray)}"
          )
  
      ndarray = jnp.asarray(ndarray)
  
      if (
          ndarray.ndim == 4 and ndarray.shape[-1] == 1
      ):  # single-channel images
          ndarray = jnp.concatenate((ndarray, ndarray, ndarray), -1)
  
      # make the mini-batch of images into a grid
      nmaps = ndarray.shape[0]
      xmaps = min(nrow, nmaps)
      ymaps = int(math.ceil(float(nmaps) / xmaps))
      height, width = (
          int(ndarray.shape[1] + padding),
          int(ndarray.shape[2] + padding),
      )
      num_channels = ndarray.shape[3]
      grid = jnp.full(
          (
              height * ymaps + padding,
              width * xmaps + padding,
              num_channels,
          ),
          pad_value,
      ).astype(jnp.float32)
      k = 0
      for y in range(ymaps):
          for x in range(xmaps):
              if k >= nmaps:
                  break
              grid = grid.at[
                  y * height + padding : (y + 1) * height,
                  x * width + padding : (x + 1) * width,
              ].set(ndarray[k])
              k = k + 1
  
      # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
      ndarr = np.array(
          jnp.clip(grid * 255.0 + 0.5, 0, 255).astype(jnp.uint8)
      )
      im = Image.fromarray(ndarr.copy())
      im.save(fp, format=format_img)
  
imports: import math
  import jax.numpy
  import numpy
  from PIL import Image
label_desc: Use image node_id 1 for calling a global function with example usage: save_image(ndarray, fp, nrow, padding, pad_value, format_img) without return types