code: def plot_images(
      ds,
      nr_rows: int,
      nr_cols: int,
      title: str,
      inference_fn: Callable | None = None,
  ):
      """Plots a grid of images with their predictions.
  
      Params:
        ds: a tensorflow dataset from where to pick the images and labels.
        nr_rows, nr_cols: the size of the grid to plot
        title: the title of the plot
        inference_fn: if None then print the existing label, else use this function
          on the batch of images to produce a batch of inference results, which
          get printed.
        inference_batch_size: the size of the batch of images passed to
        `inference_fn`.
      """
      count = nr_rows * nr_cols
      fig = plt.figure(figsize=(8.0, 4.0), num=title)
      # Get the first batch
      ((images, labels),) = list(tfds.as_numpy(ds.take(1)))
      if inference_fn:
          inferred_labels = inference_fn(images)
      for i, image in enumerate(images[:count]):
          digit = fig.add_subplot(nr_rows, nr_cols, i + 1)
          if inference_fn:
              digit_title = f"infer: {np.argmax(inferred_labels[i])}\n"
          else:
              digit_title = ""
          digit_title += f"label: {np.argmax(labels[i])}"
          digit.set_title(digit_title)
          plt.imshow(
              (np.reshape(image, (28, 28)) * 255).astype(np.uint8),
              interpolation="nearest",
          )
      plt.show()
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import logging
  import re
  import time
  from typing import Any, Callable, Optional
  from absl import flags
  import flax
  from flax import linen
  import jax
  import jax.numpy
  from matplotlib import pyplot
  import numpy
  import optax
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: plot_images(ds, nr_rows, nr_cols, title, inference_fn) without return types