diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index e295490361..54269772f1 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -39,7 +39,7 @@ async function renderDag(graph, additions, recenter=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})`);
     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("style", d => `stroke:#4a4b57; stroke-width:${STROKE_WIDTH}px; ${d.style}`);
+      .attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("style", d => d.style ?? `stroke:#4a4b57; stroke-width:${STROKE_WIDTH}px;`);
     nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label").attr("transform", d => {
       const x = (d.width-d.padding*2)/2;
       const y = (d.height-d.padding*2)/2+STROKE_WIDTH;
@@ -280,21 +280,21 @@ async function main() {
   }
   const kernelList = document.querySelector(".kernel-list");
   kernelList.innerHTML = "";
-  for (const [i,k] of kernels.entries()) {
+  for (const [i,[name, steps]] of kernels.entries()) {
     const ul = kernelList.appendChild(document.createElement("ul"));
     if (i === currentKernel) {
       ul.className = "active";
       requestAnimationFrame(() => ul.scrollIntoView({ behavior: "auto", block: "nearest" }));
     }
     const p = ul.appendChild(document.createElement("p"));
-    p.innerHTML = k[0].replace(/\u001b\[(\d+)m(.*?)\u001b\[0m/g, (_, code, st) => {
+    p.innerHTML = name.replace(/\u001b\[(\d+)m(.*?)\u001b\[0m/g, (_, code, st) => {
       const colors = ['gray','red','green','yellow','blue','magenta','cyan','white'];
       return `<span style="${`color: color-mix(in srgb, ${colors[(parseInt(code)-30+60)%60]} 60%, white)`}">${st}</span>`;
     });
     p.onclick = () => {
       setState(i === currentKernel ? { expandKernel:!expandKernel } : { expandKernel:true, currentKernel:i, currentUOp:0, currentRewrite:0 });
     }
-    for (const [j,u] of k[1].entries()) {
+    for (const [j,u] of steps.entries()) {
       const inner = ul.appendChild(document.createElement("ul"));
       if (i === currentKernel && j === currentUOp) {
         inner.className = "active";
diff --git a/tinygrad/viz/js/worker.js b/tinygrad/viz/js/worker.js
index 973a484616..3d77dd0db0 100644
--- a/tinygrad/viz/js/worker.js
+++ b/tinygrad/viz/js/worker.js
@@ -8,15 +8,16 @@ onmessage = (e) => {
   const { graph, additions } = e.data;
   const g = new dagre.graphlib.Graph({ compound: true });
   g.setGraph({ rankdir: "LR" }).setDefaultEdgeLabel(function() { return {}; });
-  if (additions.length !== 0) g.setNode("addition", {label:"", style:"fill: rgba(26, 27, 38, 0.5); stroke: none;", padding:0});
-  for (const [k, {label, src, color}] of Object.entries(graph)) {
+  if (additions.length !== 0) g.setNode("addition", {label:"", style:"fill: rgba(26, 27, 38, 0.5);", padding:0});
+  for (const [k, {label, src, ...rest }] of Object.entries(graph)) {
     // adjust node dims by label size + add padding
     let [width, height] = [0, 0];
     for (line of label.split("\n")) {
       width = Math.max(width, ctx.measureText(line).width);
       height += LINE_HEIGHT;
     }
-    g.setNode(k, {label, color, width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING});
+    g.setNode(k, {width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING, label, ...rest});
+    // add edges
     const edgeCounts = {}
     for (const s of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
     for (const s of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? edgeCounts[s] : null });