code: def _make_grid(plot_single_ex_fn, ds, rows, cols, plot_scale):
      """Plot each individual example in a grid.
  
      Args:
        plot_single_ex_fn: Function with fill a single cell of the grid, with
          signature `fn(ax: matplotlib.axes.Axes, ex: Nested[np.array]) -> None`
        ds: `tf.data.Dataset`. The tf.data.Dataset object to visualize. Examples
          should not be batched. Examples will be consumed in order until (rows *
          cols) are read or the dataset is consumed.
        rows: `int`, number of rows of the display grid.
        cols: `int`, number of columns of the display grid.
        plot_scale: `float`, controls the plot size of the images. Keep this value
          around 3 to get a good plot. High and low values may cause the labels to
          get overlapped.
  
      Returns:
        fig: The `matplotlib.Figure` object.
      """
      plt = lazy_imports_lib.lazy_imports.matplotlib.pyplot
  
      num_examples = rows * cols
      examples = list(dataset_utils.as_numpy(ds.take(num_examples)))
  
      fig = plt.figure(figsize=(plot_scale * cols, plot_scale * rows))
      fig.subplots_adjust(hspace=1 / plot_scale, wspace=1 / plot_scale)
  
      for i, ex in enumerate(examples):
          ax = fig.add_subplot(rows, cols, i + 1)
          plot_single_ex_fn(ax, ex)
  
      plt.show()
      return fig
  
file_path: datasets/tensorflow_datasets/core/visualization/image_visualizer.py
imports: from __future__ import annotations
  from typing import Optional
  from absl import logging
  from tensorflow_datasets.core import dataset_info
  from tensorflow_datasets.core import dataset_utils
  from tensorflow_datasets.core import features
  from tensorflow_datasets.core import lazy_imports_lib
  from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow
  from tensorflow_datasets.core.visualization import visualizer
label_desc: Use image node_id 1 for calling a global function with example usage: _make_grid(plot_single_ex_fn, ds, rows, cols, plot_scale) and returns: fig