Fix layer_types filtering to use layer_name for semantic categories

The graph stores semantic categories (DL_Layer, MLP_Layer, Activation, etc.)
in layer_name field, not layer_type. layer_type is always 'ATen_Operation'
for PyTorch ops.

- Add _get_node_category() helper to check both layer_name and layer_type
- Update visualize_relevance() to use _get_node_category for filtering
- Update visualize_relevance_fast() to use _get_node_category
- Update visualize_relevance_auto() node counting
- Fix color map lookup to use semantic categories

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: neerajaryaai

dl_backtrace/pytorch_backtrace/dlbacktrace/core/visualization.py

@@ -9,23 +9,48 @@
 from IPython.display import display, SVG, Image as IPyImage
 from typing import Optional, Sequence
 
-# Semantically meaningful layer types for compact visualization
+# Semantically meaningful layer categories for compact visualization
+# These match the ATEN_LAYER_MAP categories in graph_builder.py
 SEMANTIC_LAYER_TYPES: tuple[str, ...] = (
-    "MLP_Layer",      # Linear/FC layers
-    "DL_Layer",       # Conv layers (Conv1d, Conv2d, Conv3d)
+    "MLP_Layer",      # Linear/FC layers (linear, addmm)
+    "DL_Layer",       # Conv layers (conv2d, max_pool2d, etc.)
     "Activation",     # ReLU, GELU, SiLU, etc.
     "Normalization",  # BatchNorm, LayerNorm, GroupNorm
     "Attention",      # Self/Cross attention (scaled_dot_product_attention)
-    "Output",         # Final output node
-    "Placeholder",    # Input nodes (x, input_ids, etc.) - CRITICAL for graph connectivity
-    "Model_Input",    # Legacy input type (kept for compatibility)
+    "Output",         # Final output node (layer_type)
+    "Placeholder",    # Input nodes (layer_type) - CRITICAL for graph connectivity
+    "Model_Input",    # Legacy input type (layer_type)
     "NLP_Embedding",  # Embedding layers (embedding, embedding_bag)
 )
 
 # Default types to always force-include (for graph connectivity)
 DEFAULT_FORCE_INCLUDE_TYPES: tuple[str, ...] = ("Placeholder", "Model_Input", "Output")
 
 
+def _get_node_category(node_attrs: dict) -> str:
+    """Get the semantic category for a node by checking both layer_type and layer_name.
+    
+    The graph stores:
+    - layer_type: 'ATen_Operation', 'Placeholder', 'Output', 'Operation', etc.
+    - layer_name: 'DL_Layer', 'MLP_Layer', 'Activation', 'Normalization', etc.
+    
+    For filtering, we need to check layer_name first (for ATen ops), then layer_type.
+    """
+    layer_name = node_attrs.get("layer_name", "")
+    layer_type = node_attrs.get("layer_type", "Unknown")
+    
+    # For ATen operations, layer_name contains the semantic category
+    if layer_name in SEMANTIC_LAYER_TYPES:
+        return layer_name
+    
+    # For Placeholder, Output, etc., layer_type is the category
+    if layer_type in ("Placeholder", "Output", "Model_Input"):
+        return layer_type
+    
+    # Return layer_type as fallback
+    return layer_type
+
+
 def visualize_graph(graph, save_path="graph.png", *, show=True, dpi=600):
     """📊 Visualize forward execution graph with dynamic scaling (shows inline + saves)"""
     num_nodes = len(graph.nodes)
@@ -112,16 +137,16 @@ def visualize_relevance(graph, all_wt, output_path="backtrace_graph",
     force_include = {
         node.replace("/", " ").replace(":", " ")
         for node in graph.nodes
-        if graph.nodes[node].get("layer_type") in DEFAULT_FORCE_INCLUDE_TYPES
+        if _get_node_category(graph.nodes[node]) in DEFAULT_FORCE_INCLUDE_TYPES
     }
 
     if layer_types is not None:
-        # Layer-type filtering mode
+        # Layer-type filtering mode - use _get_node_category for proper semantic matching
         layer_types_set = set(layer_types)
         top_node_names = {
             node.replace("/", " ").replace(":", " ")
             for node in graph.nodes
-            if graph.nodes[node].get("layer_type") in layer_types_set
+            if _get_node_category(graph.nodes[node]) in layer_types_set
         } | force_include
     elif top_k:
         top_keys = sorted(flat_scores.items(), key=lambda x: abs(x[1]), reverse=True)[:top_k]
@@ -154,7 +179,7 @@ def find_filtered_ancestors(node_raw, visited=None):
                 ancestors.update(find_filtered_ancestors(parent_raw, visited))
         return ancestors
 
-    # --- Color map for node types ---
+    # --- Color map for node types (uses layer_name for ATen ops, layer_type for others) ---
     color_map = {
         "MLP_Layer": "lightblue",
         "DL_Layer": "lightgreen",
@@ -184,7 +209,9 @@ def find_filtered_ancestors(node_raw, visited=None):
         if name not in top_node_names:
             continue
         rel = relevance_data.get(name, (0.0, 0.0, 0.0))
-        fill = color_map.get(graph.nodes[node].get("layer_type", "Unknown"), "white")
+        # Use layer_name first (for semantic category), then layer_type as fallback
+        node_category = _get_node_category(graph.nodes[node])
+        fill = color_map.get(node_category, color_map.get(graph.nodes[node].get("layer_type", "Unknown"), "white"))
         g.node(
             name,
             label=f"{name}\nMean: {rel[0]:.3f}\nMax: {rel[1]:.3f}\nMin: {rel[2]:.3f}",
@@ -356,12 +383,12 @@ def _norm(s):
     # present nodes - keep all for transitive edge computation
     all_raw = list(graph.nodes.keys())
 
-    # Determine filtered set
+    # Determine filtered set using _get_node_category for proper semantic matching
     if layer_types is not None:
         layer_types_set = set(layer_types) | set(DEFAULT_FORCE_INCLUDE_TYPES)
         present_raw = [
             raw for raw in all_raw
-            if graph.nodes[raw].get("layer_type") in layer_types_set
+            if _get_node_category(graph.nodes[raw]) in layer_types_set
         ]
     else:
         present_raw = all_raw
@@ -479,8 +506,9 @@ def _short(s, n=48):
     for raw in present_raw:
         nk = norm_by_raw[raw]
         mean, mx, mn = rel_map.get(nk, (0.0, 0.0, 0.0))
-        lt = graph.nodes[raw].get("layer_type", "Unknown")
-        fill = color_map.get(lt, "white")
+        # Use _get_node_category for proper semantic coloring
+        node_category = _get_node_category(graph.nodes[raw])
+        fill = color_map.get(node_category, color_map.get(graph.nodes[raw].get("layer_type", "Unknown"), "white"))
         collapsed = graph.nodes[raw].get("collapsed_count", 0)
         collapsed_line = f"\n[collapsed {collapsed}]" if collapsed else ""
 
@@ -568,7 +596,7 @@ def visualize_relevance_auto(
         layer_types_set = set(layer_types) | set(DEFAULT_FORCE_INCLUDE_TYPES)
         filtered_count = sum(
             1 for n in graph.nodes 
-            if graph.nodes[n].get("layer_type") in layer_types_set
+            if _get_node_category(graph.nodes[n]) in layer_types_set
         )
         num_nodes = filtered_count
         print(f"num_nodes after layer_types filter: {num_nodes} (from {len(graph.nodes)} total)")