From 134d92a1703e40059744824f653a7016a6968516 Mon Sep 17 00:00:00 2001
From: Nhat Minh Luu <54903938+nluu175@users.noreply.github.com>
Date: Tue, 29 Oct 2024 21:41:42 -0600
Subject: [PATCH 1/4] Modify example code to fix torch logs doc

---
 recipes_source/torch_logs.py | 89 +++++++++++++++++++++---------------
 1 file changed, 51 insertions(+), 38 deletions(-)

diff --git a/recipes_source/torch_logs.py b/recipes_source/torch_logs.py
index b5c3f0bd8a..3d6c094293 100644
--- a/recipes_source/torch_logs.py
+++ b/recipes_source/torch_logs.py
@@ -31,52 +31,65 @@
 #   variable setting is shown for each example.
 
 import torch
+import sys
 
-# exit cleanly if we are on a device that doesn't support torch.compile
-if torch.cuda.get_device_capability() < (7, 0):
-    print("Skipping because torch.compile is not supported on this device.")
-else:
-    @torch.compile()
-    def fn(x, y):
-        z = x + y
-        return z + 2
-
-
-    inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
 
+def env_setup():
+    """Set up environment for running the example. Exit cleanly if CUDA is not available."""
+    if not torch.cuda.is_available():
+        print("CUDA is not available. Exiting.")
+        sys.exit(0)
+        
+    if torch.cuda.get_device_capability() < (7, 0):
+        print("Skipping because torch.compile is not supported on this device.")
+        sys.exit(0)
 
-# print separator and reset dynamo
-# between each example
-    def separator(name):
-        print(f"==================={name}=========================")
-        torch._dynamo.reset()
 
+def separator(name):
+    """Print separator and reset dynamo between each example"""
+    print(f"\n{'='*20} {name} {'='*20}")
+    torch._dynamo.reset()
 
-    separator("Dynamo Tracing")
-# View dynamo tracing
-# TORCH_LOGS="+dynamo"
-    torch._logging.set_logs(dynamo=logging.DEBUG)
-    fn(*inputs)
 
-    separator("Traced Graph")
-# View traced graph
-# TORCH_LOGS="graph"
-    torch._logging.set_logs(graph=True)
-    fn(*inputs)
+def run_debugging_suite():
+    """Run the complete debugging suite with all logging options"""
+    env_setup()
 
-    separator("Fusion Decisions")
-# View fusion decisions
-# TORCH_LOGS="fusion"
-    torch._logging.set_logs(fusion=True)
-    fn(*inputs)
-
-    separator("Output Code")
-# View output code generated by inductor
-# TORCH_LOGS="output_code"
-    torch._logging.set_logs(output_code=True)
-    fn(*inputs)
+    @torch.compile()
+    def fn(x, y):
+        z = x + y
+        return z + 2
 
-    separator("")
+    inputs = (
+        torch.ones(2, 2, device="cuda"),
+        torch.zeros(2, 2, device="cuda")
+    )
+    
+    logging_scenarios = [
+        # View dynamo tracing; TORCH_LOGS="+dynamo"
+        ("Dynamo Tracing", {"dynamo": logging.DEBUG}), 
+
+        # View traced graph; TORCH_LOGS="graph"
+        ("Traced Graph", {"graph": True}), 
+        
+        # View fusion decisions; TORCH_LOGS="fusion"
+        ("Fusion Decisions", {"fusion": True}), 
+        
+        # View output code generated by inductor; TORCH_LOGS="output_code"
+        ("Output Code", {"output_code": True}) 
+    ]
+    
+    for name, log_config in logging_scenarios:
+        separator(name)
+        torch._logging.set_logs(**log_config)
+        try:
+            result = fn(*inputs)
+            print(f"Function output shape: {result.shape}")
+        except Exception as e:
+            print(f"Error during {name}: {str(e)}")
+
+if __name__ == "__main__":
+    run_debugging_suite()
 
 ######################################################################
 # Conclusion

From 368e0624db08b2e4a0519719c531fc1473920835 Mon Sep 17 00:00:00 2001
From: Nhat Minh Luu <54903938+nluu175@users.noreply.github.com>
Date: Thu, 7 Nov 2024 01:14:25 -0700
Subject: [PATCH 2/4] Introduced TORCH_TRACE/tlparse

---
 recipes_source/torch_logs.py | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/recipes_source/torch_logs.py b/recipes_source/torch_logs.py
index 3d6c094293..df1f9dc81d 100644
--- a/recipes_source/torch_logs.py
+++ b/recipes_source/torch_logs.py
@@ -88,8 +88,24 @@ def fn(x, y):
         except Exception as e:
             print(f"Error during {name}: {str(e)}")
 
-if __name__ == "__main__":
-    run_debugging_suite()
+
+run_debugging_suite()
+
+######################################################################
+# Use TORCH_TRACE/tlparse to produce produce compilation reports
+# ~~~~~~~~~~
+#
+# In this section, we introduce the usage of TORCH_TRACE and tlparse to produce reports.
+# First, we run `TORCH_TRACE="/tmp/tracedir" python script.py` to generate the raw trace logs.
+# We have replace `/tmp/tracedir` with a path to a directory you want to store the trace logs
+# and reaplce script with the name of your script.
+# 
+# Next, we are going to pass the trace log to `tlparse` to generate compilation reports. We run
+# `pip install tlparse` and then `tlparse /tmp/tracedir`. This will open up your browser with 
+# HTML like generated above.
+#
+# By default, reports generated by `tlparse` are 
+# stored in the directory `tl_out`. You can change that by doing `tlparse /tmp/tracedir -o output_dir/`.
 
 ######################################################################
 # Conclusion

From 8859a58aad42fbd9b50b28707665f63775dc6f0b Mon Sep 17 00:00:00 2001
From: Nhat Minh Luu <54903938+nluu175@users.noreply.github.com>
Date: Thu, 7 Nov 2024 13:56:27 -0700
Subject: [PATCH 3/4] Adjust the wording based on suggestions

---
 recipes_source/torch_logs.py | 47 ++++++++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 13 deletions(-)

diff --git a/recipes_source/torch_logs.py b/recipes_source/torch_logs.py
index df1f9dc81d..9a0e959ce9 100644
--- a/recipes_source/torch_logs.py
+++ b/recipes_source/torch_logs.py
@@ -35,7 +35,7 @@
 
 
 def env_setup():
-    """Set up environment for running the example. Exit cleanly if CUDA is not available."""
+    """Set up the environment to run the example. Exit cleanly if CUDA is not available."""
     if not torch.cuda.is_available():
         print("CUDA is not available. Exiting.")
         sys.exit(0)
@@ -46,7 +46,7 @@ def env_setup():
 
 
 def separator(name):
-    """Print separator and reset dynamo between each example"""
+    """Print a separator and reset dynamo between each example"""
     print(f"\n{'='*20} {name} {'='*20}")
     torch._dynamo.reset()
 
@@ -92,20 +92,41 @@ def fn(x, y):
 run_debugging_suite()
 
 ######################################################################
-# Use TORCH_TRACE/tlparse to produce produce compilation reports
-# ~~~~~~~~~~
+# Using ``TORCH_TRACE/tlparse`` to produce produce compilation reports (for PyTorch 2)
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# In this section, we introduce ``TORCH_TRACE`` and ``tlparse`` to produce reports.
+# 
+# 
+# 1. Generate the raw trace logs by running the following command:
+# 
+# .. code-block:: bash
+# 
+#    TORCH_TRACE="/tmp/tracedir" python script.py`
+#
+# Ensure you replace ``/tmp/tracedir`` with the path to the directory where you want 
+# to store the trace logs and replace the script with the name of your script.
+# 
+# 2. Install ``tlparse`` by running:
+#
+#    .. code-block:: bash
+#
+#       pip install tlparse
+#
+# 3. Pass the trace log to ``tlparse`` to generate compilation reports:
+#
+#     .. code-block: bash
+#
+#        tlparse /tmp/tracedir
 #
-# In this section, we introduce the usage of TORCH_TRACE and tlparse to produce reports.
-# First, we run `TORCH_TRACE="/tmp/tracedir" python script.py` to generate the raw trace logs.
-# We have replace `/tmp/tracedir` with a path to a directory you want to store the trace logs
-# and reaplce script with the name of your script.
+# This will open your browser with the HTML-like generated above.  
 # 
-# Next, we are going to pass the trace log to `tlparse` to generate compilation reports. We run
-# `pip install tlparse` and then `tlparse /tmp/tracedir`. This will open up your browser with 
-# HTML like generated above.
+# By default, reports generated by ``tlparse`` are stored in the ``tl_out`` directory. 
+# You can change that by running:
+#  
+# .. code-block:: bash
 #
-# By default, reports generated by `tlparse` are 
-# stored in the directory `tl_out`. You can change that by doing `tlparse /tmp/tracedir -o output_dir/`.
+#    tlparse /tmp/tracedir -o output_dir/
 
 ######################################################################
 # Conclusion

From 26d2b0e7873a9d903360242525c3f44a438f95d6 Mon Sep 17 00:00:00 2001
From: Svetlana Karslioglu <svekars@meta.com>
Date: Thu, 7 Nov 2024 17:20:36 -0800
Subject: [PATCH 4/4] Update recipes_source/torch_logs.py

---
 recipes_source/torch_logs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes_source/torch_logs.py b/recipes_source/torch_logs.py
index 9a0e959ce9..948ffd5571 100644
--- a/recipes_source/torch_logs.py
+++ b/recipes_source/torch_logs.py
@@ -119,7 +119,7 @@ def fn(x, y):
 #
 #        tlparse /tmp/tracedir
 #
-# This will open your browser with the HTML-like generated above.  
+# This will open your browser with the HTML-like code generated above.  
 # 
 # By default, reports generated by ``tlparse`` are stored in the ``tl_out`` directory. 
 # You can change that by running: