def _extract_image_batch(images, num_images, batch_size):
    """Extracts a batch of images for plotting.

    Args:
        images: The 4D tensor or NumPy array of images.
        num_images: The number of images to extract.
        batch_size: The original batch size of the images.

    Returns:
        A 4D tensor or NumPy array containing the extracted images.

    Raises:
        ValueError: If `images` is not a 4D tensor/array.
    """

    if len(ops.shape(images)) != 4:
        raise ValueError(
            "`plot_images_gallery()` requires you to "
            "batch your `np.array` samples together."
        )
    num_samples = min(num_images, batch_size)
    sample = images[:num_samples, ...]

    return sample
Extracts a batch of images for plotting.

    Args:
        images: The 4D tensor or NumPy array of images.
        num_images: The number of images to extract.
        batch_size: The original batch size of the images.

    Returns:
        A 4D tensor or NumPy array containing the extracted images.

    Raises:
        ValueError: If `images` is not a 4D tensor/array.
    
keras/src/visualization/plot_image_gallery.py