Improve graph visualization documentation

docs/source/_templates/function.rst

@@ -0,0 +1,8 @@
+.. role:: hidden
+ :class: hidden-section
+.. currentmodule:: {{ module }}
+
+
+{{ name | underline }}
+
+.. autofunction:: {{ name }}

docs/source/torchdata.datapipes.utils.rst

@@ -1,6 +1,21 @@
 Utility Functions
 ===========================
 
+..
+ Comment: the next section will become "DataPipe Graph Visualization and Linter" when linter is added.
+
+DataPipe Graph Visualization
+-------------------------------------
+.. currentmodule:: torchdata.datapipes.utils
+
+.. autosummary::
+ :nosignatures:
+ :toctree: generated/
+ :template: function.rst
+
+ to_graph
+
+
 File Object and Stream Utility
 -------------------------------------
 
@@ -12,7 +27,6 @@ File Object and Stream Utility
  :template: datapipe.rst
 
  StreamWrapper
- to_graph
 
 
 DataLoader

torchdata/datapipes/utils/_visualization.py

@@ -115,7 +115,8 @@ def aggregate(nodes):
 
 
 def to_graph(dp, *, debug: bool = False) -> "graphviz.Digraph":
- """Turns a datapipe into a :class:`graphviz.Digraph` representing the graph of the datapipe.
+ """Visualizes a DataPipe by returning a :class:`graphviz.Digraph`, which is a graph of the data pipeline.
+ This allows you to visually inspect all the transformation that takes place in your DataPipes.
 
  .. note::
 
@@ -129,9 +130,20 @@ def to_graph(dp, *, debug: bool = False) -> "graphviz.Digraph":
  - :meth:`~graphviz.Digraph.view`: Open the graph in a viewer.
 
  Args:
- dp: Datapipe.
- debug (bool): If ``True``, renders internal datapipes that are usually hidden from the user. Defaults to
- ``False``.
+ dp: DataPipe that you would like to visualize (generally the last one in a chain of DataPipes).
+ debug (bool): If ``True``, renders internal datapipes that are usually hidden from the user
+ (such as ``ChildDataPipe`` of `demux` and `fork`). Defaults to ``False``.
+
+ Example:
+ >>> from torchdata.datapipes.iter import IterableWrapper
+ >>> from torchdata.datapipes.utils import to_graph
+ >>> dp = IterableWrapper(range(10))
+ >>> dp1, dp2 = dp.demux(num_instances=2, classifier_fn=lambda x: x % 2)
+ >>> dp1 = dp1.map(lambda x: x + 1)
+ >>> dp2 = dp2.filter(lambda _: True)
+ >>> dp3 = dp1.zip(dp2).map(lambda t: t[0] + t[1])
+ >>> g = to_graph(dp3)
+ >>> g.view() # This will open the graph in a viewer
  """
  try:
  import graphviz