From 901d27b3ba8549e845abe975e3fdff3cf8751c92 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Tue, 28 Oct 2025 18:54:28 +0800
Subject: [PATCH] viz: optional text dims try 2 (#12971)

---
 tinygrad/viz/js/index.js | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index cceae08cea..3f72816cc9 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -70,11 +70,10 @@ const drawGraph = (data) => {
   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);
   const STROKE_WIDTH = 1.4;
-  nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label").attr("transform", d => {
-    const x = d.labelWidth/2;
-    const y = d.labelHeight/2+STROKE_WIDTH*2;
-    return `translate(-${x}, -${y})`;
-  }).selectAll("text").data(d => {
+  const labels = nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label");
+  const hasLabelDims = data.nodes[0]?.value.labelWidth != null;
+  if (hasLabelDims) labels.attr("transform", d => `translate(-${d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`);
+  labels.selectAll("text").data(d => {
     const ret = [[]];
     for (const { st, color } of parseColors(d.label, defaultColor="initial")) {
       const lines = st.split("\n");
@@ -84,6 +83,11 @@ const drawGraph = (data) => {
     return [ret];
   }).join("text").selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", 14).selectAll("tspan").data(d => d).join("tspan")
     .attr("fill", d => darkenHex(d.color, 25)).text(d => d.st).attr("xml:space", "preserve");
+  // recenter after drawing texts if needed
+  if (!hasLabelDims) labels.attr("transform", (_,i,els) => {
+    const b = els[i].getBBox();
+    return `translate(${-b.x-b.width/2}, ${-b.y-b.height/2})`
+  });
   addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
     .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
   // draw edges