From 4cd27aa0e6919f877e141a354823113fbd179bb3 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Mon, 7 Apr 2025 14:38:03 +0800
Subject: [PATCH] hotfix: viz recenter and unlimited zoom (#9760)

* hotfix: viz recenter and unlimited zoom

* add shapes to the ast graph

* not for COPY
---
 tinygrad/viz/js/index.js | 3 ++-
 tinygrad/viz/serve.py    | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index d22a6dd208..f395ebcc59 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -190,7 +190,7 @@ function renderMemoryGraph(graph) {
 
 // ** zoom and recentering
 
-const zoom = d3.zoom().scaleExtent([0.05, 2]).on("zoom", (e) => d3.select("#render").attr("transform", e.transform));
+const zoom = d3.zoom().on("zoom", (e) => d3.select("#render").attr("transform", e.transform));
 d3.select("#graph-svg").call(zoom);
 // zoom to fit into view
 document.getElementById("zoom-to-fit-btn").addEventListener("click", () => {
@@ -426,6 +426,7 @@ document.addEventListener("keydown", async function(event) {
     return setState({ currentRewrite:Math.min(totalRewrites, currentRewrite+1) });
   }
   if (event.key == " ") {
+    event.preventDefault()
     document.getElementById("zoom-to-fit-btn").click();
   }
 });
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index 308490fa30..96f5aa9b51 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -71,6 +71,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
       if x in excluded:
         if x.op is Ops.CONST and dtypes.is_float(u.dtype): label += f"\nCONST{idx} {x.arg:g}"
         else: label += f"\n{x.op.name}{idx} {x.arg}"
+    if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, *GroupOp.Buffer} and u.st is not None: label += f"\n{repr(u.shape)}"
     graph[id(u)] = {"label":label, "src":[id(x) for x in u.src if x not in excluded], "color":uops_colors.get(u.op, "#ffffff")}
   return graph