From 3674c0754ecd2fae4cf8470d83d69fa486d1fd40 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Tue, 26 Aug 2025 14:56:13 +0300
Subject: [PATCH] viz: small uop click changes (#11846)

* also highlight self

* can always unselect by clicking outside

* less layout
---
 tinygrad/viz/index.html  | 3 ---
 tinygrad/viz/js/index.js | 6 ++++--
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html
index 0807fbbab6..aec595d30f 100644
--- a/tinygrad/viz/index.html
+++ b/tinygrad/viz/index.html
@@ -100,9 +100,6 @@
   .highlight rect, .edgePath.highlight {
     stroke: #FFC53D;
   }
-  .highlight rect {
-    stroke-width: 3px;
-  }
   #arrowhead {
     fill: #4a4b57;
   }
diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index bab878708f..42bdba1166 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -56,12 +56,14 @@ async function renderDag(graph, additions, recenter=false) {
     const g = dagre.graphlib.json.read(e.data);
     // draw nodes
     const STROKE_WIDTH = 1.4;
+    d3.select("#graph-svg").on("click", () => d3.selectAll(".highlight").classed("highlight", false));
     const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g")
-      .attr("transform", d => `translate(${d.x},${d.y})`).classed("clickable", d => d.ref != null).on("click", (_,d) => {
+      .attr("transform", d => `translate(${d.x},${d.y})`).classed("clickable", d => d.ref != null).on("click", (e,d) => {
         if (d.ref != null) return setCtxWithHistory(d.ref);
-        const src = g.predecessors(d.id) || [];
+        const src = [...g.predecessors(d.id), d.id];
         nodes.classed("highlight", n => src.includes(n.id));
         d3.select("#edges").selectAll("path.edgePath").classed("highlight", e => src.includes(e.v) && e.w===d.id);
+        e.stopPropagation();
       });
     nodes.selectAll("rect").data(d => [d]).join("rect").attr("width", d => d.width).attr("height", d => d.height).attr("fill", d => d.color)
       .attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("class", d => d.className ?? "node");