3535
3636
3737def env_setup ():
38- """Set up environment for running the example. Exit cleanly if CUDA is not available."""
38+ """Set up the environment to run the example. Exit cleanly if CUDA is not available."""
3939 if not torch .cuda .is_available ():
4040 print ("CUDA is not available. Exiting." )
4141 sys .exit (0 )
@@ -46,7 +46,7 @@ def env_setup():
4646
4747
4848def separator (name ):
49- """Print separator and reset dynamo between each example"""
49+ """Print a separator and reset dynamo between each example"""
5050 print (f"\n { '=' * 20 } { name } { '=' * 20 } " )
5151 torch ._dynamo .reset ()
5252
@@ -92,20 +92,41 @@ def fn(x, y):
9292run_debugging_suite ()
9393
9494######################################################################
95- # Use TORCH_TRACE/tlparse to produce produce compilation reports
96- # ~~~~~~~~~~
95+ # Using ``TORCH_TRACE/tlparse`` to produce produce compilation reports (for PyTorch 2)
96+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
97+ #
98+ # In this section, we introduce ``TORCH_TRACE`` and ``tlparse`` to produce reports.
99+ #
100+ #
101+ # 1. Generate the raw trace logs by running the following command:
102+ #
103+ # .. code-block:: bash
104+ #
105+ # TORCH_TRACE="/tmp/tracedir" python script.py`
106+ #
107+ # Ensure you replace ``/tmp/tracedir`` with the path to the directory where you want
108+ # to store the trace logs and replace the script with the name of your script.
109+ #
110+ # 2. Install ``tlparse`` by running:
111+ #
112+ # .. code-block:: bash
113+ #
114+ # pip install tlparse
115+ #
116+ # 3. Pass the trace log to ``tlparse`` to generate compilation reports:
117+ #
118+ # .. code-block: bash
119+ #
120+ # tlparse /tmp/tracedir
97121#
98- # In this section, we introduce the usage of TORCH_TRACE and tlparse to produce reports.
99- # First, we run `TORCH_TRACE="/tmp/tracedir" python script.py` to generate the raw trace logs.
100- # We have replace `/tmp/tracedir` with a path to a directory you want to store the trace logs
101- # and reaplce script with the name of your script.
122+ # This will open your browser with the HTML-like generated above.
102123#
103- # Next, we are going to pass the trace log to `tlparse` to generate compilation reports. We run
104- # `pip install tlparse` and then `tlparse /tmp/tracedir`. This will open up your browser with
105- # HTML like generated above.
124+ # By default, reports generated by ``tlparse`` are stored in the ``tl_out`` directory.
125+ # You can change that by running:
126+ #
127+ # .. code-block:: bash
106128#
107- # By default, reports generated by `tlparse` are
108- # stored in the directory `tl_out`. You can change that by doing `tlparse /tmp/tracedir -o output_dir/`.
129+ # tlparse /tmp/tracedir -o output_dir/
109130
110131######################################################################
111132# Conclusion
0 commit comments