From fef97547f9de4c3b2f79f84ce7a3c0474f43de84 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Mon, 18 Aug 2025 17:51:21 +0300
Subject: [PATCH] viz: preset the final timestamp (#11712)

---
 tinygrad/viz/js/index.js | 2 +-
 tinygrad/viz/serve.py    | 8 +++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index d1afdec45a..b8895a2179 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -126,7 +126,7 @@ const createPolygons = (source, height) => {
   const shapes = [];
   const yscale = d3.scaleLinear().domain([0, source.peak]).range([height, 0]);
   for (const [i,e] of source.shapes.entries()) {
-    const x = e.x.map((i,_) => (source.timestamps[i] ?? data.et)-data.st);
+    const x = e.x.map((i,_) => source.timestamps[i]-data.st);
     const y0 = e.y.map(yscale);
     const y1 = e.y.map(y => yscale(y+e.arg.nbytes));
     const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` };
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index e9ad038c9e..f03c2894f9 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -144,7 +144,7 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
     shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info})
   return {"shapes":shapes, "maxDepth":len(levels)}
 
-def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dict:
   step, peak, mem = 0, 0, 0
   shps:dict[int, dict] = {}
   temp:dict[int, dict] = {}
@@ -170,9 +170,10 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
   for v in temp.values():
     v["x"].append(step)
     v["y"].append(v["y"][-1])
+  timestamps.append(max_ts)
   return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}
 
-def get_profile(profile:list[ProfileEvent]):
+def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   # start by getting the time diffs
   for ev in profile:
     if isinstance(ev,ProfileDeviceEvent): device_ts_diffs[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
@@ -184,12 +185,13 @@ def get_profile(profile:list[ProfileEvent]):
     dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en), float(en-ts), e))
     if min_ts is None or st < min_ts: min_ts = st
     if max_ts is None or et > max_ts: max_ts = et
+  if min_ts is None: return None
   # return layout of per device events
   layout:dict[str, dict] = {}
   for k,v in dev_events.items():
     v.sort(key=lambda e:e[0])
     layout[k] = timeline_layout(v)
-    layout[f"{k} Memory"] = mem_layout(v)
+    layout[f"{k} Memory"] = mem_layout(v, unwrap(max_ts))
   return json.dumps({"layout":layout, "st":min_ts, "et":max_ts}).encode("utf-8")
 
 def get_runtime_stats(key) -> list[dict]: