Added Pt Arange and Pt Linspace nodes.

dev/scripts/generate_node_breakdown_csv.py

@@ -0,0 +1,63 @@
+import re
+from typing import Dict
+import pandas as pd
+
+def increment_count(d: Dict[str, int], category: str) -> None:
+    """Increment the count of a category in the dictionary.
+
+    Args:
+        d (Dict[str, int]): Dictionary storing category counts.
+        category (str): The category to increment.
+    """
+    d[category] = d.get(category, 0) + 1
+
+def process_import_lines(file_path: str) -> Dict[str, int]:
+    """Extract and categorize import lines from a Python file.
+
+    Args:
+        file_path (str): Path to the Python file.
+
+    Returns:
+        Dict[str, int]: Dictionary containing category counts.
+    """
+    category_counts = {}
+    
+    with open(file_path, "r") as f:
+        for line in f:
+            if "import" in line:
+                module_name = re.sub(r"^.*import ", "", line).strip()
+                if module_name.startswith("SNS"):
+                    increment_count(category_counts, "Seaborn")
+                elif module_name.startswith("MPL"):
+                    increment_count(category_counts, "Matplotlib")
+                elif module_name.startswith("Pt"):
+                    increment_count(category_counts, "PyTorch")
+                elif module_name.startswith("Pandas"):
+                    increment_count(category_counts, "Pandas")
+                elif module_name.startswith(("CDA", "Python")):
+                    increment_count(category_counts, "Misc")
+                elif module_name.startswith("Numpy"):
+                    increment_count(category_counts, "NumPy")
+                else:
+                    increment_count(category_counts, "Misc")
+    
+    return category_counts
+
+def save_category_counts_to_csv(category_counts: Dict[str, int], output_file: str) -> None:
+    """Save category counts as a sorted CSV file.
+
+    Args:
+        category_counts (Dict[str, int]): Dictionary containing category counts.
+        output_file (str): Path to the output CSV file.
+    """
+    df = pd.DataFrame(list(category_counts.items()), columns=["Category", "Count"])
+    df = df.sort_values(by=["Count"], ascending=False)
+    df.to_csv(output_file, index=False)
+
+def main() -> None:
+    """Main function to process the import categories and save results."""
+    category_counts = process_import_lines("__init__.py")
+    save_category_counts_to_csv(category_counts, "node_count_breakdown.csv")
+
+if __name__ == "__main__":
+    main()