def METHOD_NAME(images, num_images, batch_size):
    def unpack_images(inputs):
        return inputs["image"]
    num_batches_required = math.ceil(num_images / batch_size)
    if isinstance(images, tf.data.Dataset):
        images = images.map(unpack_images)
        if batch_size == 1:
            images = images.ragged_batch(num_batches_required)
            sample = next(iter(images.take(1)))
        else:
            sample = next(iter(images.take(num_batches_required)))
        return sample
    else:
        if len(images.shape) != 4:
            raise ValueError(
                "`plot_images_gallery()` requires you to "
                "batch your `np.array` samples together."
            )
        else:
            num_samples = (
                num_images if num_images <= batch_size else num_batches_required
            )
            sample = images[:num_samples, ...]
    return sample