From 1858f1fd9aa94ca4e302b60a88f075d0d1dd88bc Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Thu, 6 Aug 2026 19:44:34 -0700
Subject: [PATCH] viz: collapse PROGRAM nodes like CALL (codex) (#17438)

Co-authored-by: qazal <77887910+Qazalin@users.noreply.github.com>
---
 tinygrad/viz/js/index.js  |  6 +++---
 tinygrad/viz/js/worker.js | 14 ++++++++------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index 0d1293f0d6..8ea2ef384a 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -120,9 +120,9 @@ const drawGraph = (data) => {
     .attr("transform", d => `translate(${d.width/2-8}, ${-d.height/2+8})`).datum(e => ({ rect:true, width:10, height:10, fill:e.addrspace, stroke:"none" })));
   const CALL_TAG_WIDTH = 14;
   addTags(nodes.selectAll("g.type").data(d => d.collapsible ? [d] : []).join("g").attr("class", d => `tag clickable ${d.collapsed ? 'collapsed' : 'expanded'}`)
-    .attr("transform", d => d.callNode ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`)
-    .datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.callNode ? null : d.color,
-      ...(d.callNode && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => {
+    .attr("transform", d => d.collapsePorts != null ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`)
+    .datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.collapsePorts != null ? null : d.color,
+      ...(d.collapsePorts != null && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => {
       e.stopPropagation();
       const t = d3.zoomTransform(document.getElementById("graph-svg"));
       const [x, y] = t.apply([d.x, d.y]);
diff --git a/tinygrad/viz/js/worker.js b/tinygrad/viz/js/worker.js
index 5a73d1b64e..00de46c5ff 100644
--- a/tinygrad/viz/js/worker.js
+++ b/tinygrad/viz/js/worker.js
@@ -54,15 +54,17 @@ const layoutUOp = (g, { graph, change }, opts) => {
       width = Math.max(width, ctx.measureText(line).width);
       height += lineHeight;
     }
-    const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n");
+    const op = label.split("\n", 1)[0];
+    const callNode = op === "CALL" || op === "FUNCTION", programNode = op === "PROGRAM";
+    const collapsePorts = callNode ? [0] : programNode ? [0, 1] : null;
     if (callNode) callCount++;
-    g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, exclude, addrspace,
+    g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, collapsePorts, exclude, addrspace,
       className:label.startsWith("REWRITE_ERROR") ? "err" : null});
     // add edges
     const edgeCounts = {};
     for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
     for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port},
-      ...(callNode && port === 0 && {color:"#a0a1b8"})});
+      ...(collapsePorts?.includes(port) && {color:"#a0a1b8"})});
     if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
   }
   // optionally hide nodes from the layout
@@ -87,11 +89,11 @@ const layoutUOp = (g, { graph, change }, opts) => {
       const consumer = g.node(consumerId);
       // add +- toggle if this consumer has collapsible sources
       const edge = g.edge(n, consumerId);
-      const collapsible = consumer.callNode ? edge?.label?.text === 0 : node.exclude;
+      const collapsible = consumer.collapsePorts != null ? consumer.collapsePorts.includes(edge?.label?.text) : node.exclude;
       if (!collapsible) continue;
       consumer.collapsible = true;
-      // increase width of call/function nodes to make space for a toggle
-      if (consumer.callNode) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; }
+      // increase width of call/function/program nodes to make space for a toggle
+      if (consumer.collapsePorts != null) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; }
       // make sources invisible if UI has toggled it off
       const collapsed = consumer.callNode ? opts.showCallSrc === opts.callSrcMask.has(consumerId) : !opts.expandedNodes.has(consumerId);
       if (!collapsed) continue;