class_docstr: Visualizer for graph prediction datasets.
class_docstr_tok: ["Visualizer","for","graph","prediction","datasets","."]
docstr: Display the dataset.
  
  Args:
    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.
    ds_info: `tfds.core.DatasetInfo` object of the dataset to visualize.
    node_label_fn: A callable that maps individual graph examples to a
      dictionary of node labels, rendered within the nodes.
    node_color_fn: A callable that maps individual graph examples to a
      dictionary of node colors.
    edge_color_fn: A callable that maps individual graph examples to a
      dictionary of edge colors.
    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 5 to get a good plot. High and low values may cause the labels to
      get overlapped.
    **kwargs: Additional arguments passed to networkx.draw_networkx().
  
  Returns:
    fig: The pyplot figure.
docstr_tok: ["Display","the","dataset",".","Args",":","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",".","ds_info",":","`","tfds.core.DatasetInfo","`","object","of","the","dataset","to","visualize",".","node_label_fn",":","A","callable","that","maps","individual","graph","examples","to","a","dictionary","of","node","labels",",","rendered","within","the","nodes",".","node_color_fn",":","A","callable","that","maps","individual","graph","examples","to","a","dictionary","of","node","colors",".","edge_color_fn",":","A","callable","that","maps","individual","graph","examples","to","a","dictionary","of","edge","colors",".","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","5","to","get","a","good","plot",".","High","and","low","values","may","cause","the","labels","to","get","overlapped",".","*","*","kwargs",":","Additional","arguments","passed","to","networkx.draw_networkx","(",")",".","Returns",":","fig",":","The","pyplot","figure","."]
code: def show(
      self,
      ds: tf.data.Dataset,
      ds_info: dataset_info.DatasetInfo,
      node_label_fn: Optional[_GraphFn] = None,
      node_color_fn: Optional[_GraphFn] = None,
      edge_color_fn: Optional[_GraphFn] = None,
      rows: int = 3,
      cols: int = 3,
      plot_scale: float = 5,
      **kwargs,
  ):
      """Display the dataset.
  
      Args:
        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.
        ds_info: `tfds.core.DatasetInfo` object of the dataset to visualize.
        node_label_fn: A callable that maps individual graph examples to a
          dictionary of node labels, rendered within the nodes.
        node_color_fn: A callable that maps individual graph examples to a
          dictionary of node colors.
        edge_color_fn: A callable that maps individual graph examples to a
          dictionary of edge colors.
        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 5 to get a good plot. High and low values may cause the labels to
          get overlapped.
        **kwargs: Additional arguments passed to networkx.draw_networkx().
  
      Returns:
        fig: The pyplot figure.
      """
  
      plt = lazy_imports_lib.lazy_imports.matplotlib.pyplot
      nx = lazy_imports_lib.lazy_imports.networkx
  
      # Construct edge-getter function which will extract edges from each graph.
      get_edges_fn = _make_edge_getter_fn(ds_info)
  
      # Create subplots.
      fig, axs = plt.subplots(
          nrows=rows,
          ncols=cols,
          squeeze=False,
          figsize=(plot_scale * cols, plot_scale * rows),
      )
      plt.subplots_adjust(hspace=1 / plot_scale, wspace=1 / plot_scale)
  
      for graph, ax in zip(ds, axs.reshape(-1)):
          # Obtain edges.
          edges = get_edges_fn(graph)
  
          # Create NetworkX graph.
          nx_graph = nx.Graph()
          nx_graph.add_edges_from(edges)
  
          # Add node labels, node colors and edge colors, if specified.
          if node_label_fn is not None:
              kwargs["labels"] = node_label_fn(graph)
  
          if node_color_fn is not None:
              node_color_map = node_color_fn(graph)
              node_color_mapper = _make_node_color_mapper(
                  node_color_map
              )
              kwargs["node_color"] = list(
                  map(node_color_mapper, nx_graph.nodes())
              )
  
          if edge_color_fn is not None:
              edge_color_map = edge_color_fn(graph)
              edge_color_mapper = _make_edge_color_mapper(
                  edge_color_map
              )
              kwargs["edge_color"] = list(
                  map(edge_color_mapper, nx_graph.edges())
              )
  
          # Draw with NetworkX.
          nx.draw_networkx(nx_graph, ax=ax, **kwargs)
  
      return fig
  
file_path: datasets/tensorflow_datasets/core/visualization/graph_visualizer.py
imports: from __future__ import annotations
  from typing import Any, Callable, Dict, Optional, Union
  from tensorflow_datasets.core import dataset_info
  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 2 for calling the GraphVisualizer obj's underlying member method code with example usage: obj.show(ds, ds_info, node_label_fn, node_color_fn, edge_color_fn, rows, cols, plot_scale) and returns: fig