diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py index 9af29a5a54..53ed18c7d6 100644 --- a/test/unit/test_viz.py +++ b/test/unit/test_viz.py @@ -274,18 +274,18 @@ def load_profile(lst:list[ProfileEvent]) -> dict: klen = u(" NV:1') self.assertEqual(nv1_events[0]['st'], 954) #self.assertEqual(j['devEvents'][7]['pid'], j['devEvents'][3]['pid']) - graph_events = j['layout']['NV Graph']['shapes'] + graph_events = j['layout']['NV Graph']['events'] self.assertEqual(graph_events[0]['st'], nv_events[0]['st']) self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur']) @@ -377,7 +377,7 @@ class TestVizMemoryLayout(BaseTestViz): profile_ret = load_profile(Buffer.profile_events) ret = profile_ret["layout"][f"{a.device} Memory"] self.assertEqual(ret["peak"], 2) - self.assertEqual(len(ret["shapes"]), 2) + self.assertEqual(len(ret["events"]), 2) def test_del_once(self): a = _alloc(1) @@ -386,7 +386,7 @@ class TestVizMemoryLayout(BaseTestViz): profile_ret = load_profile(Buffer.profile_events) ret = profile_ret["layout"][f"{b.device} Memory"] self.assertEqual(ret["peak"], 1) - self.assertEqual(len(ret["shapes"]), 3) + self.assertEqual(len(ret["events"]), 3) def test_alloc_free(self): a = _alloc(1) @@ -396,7 +396,7 @@ class TestVizMemoryLayout(BaseTestViz): profile_ret = load_profile(Buffer.profile_events) ret = profile_ret["layout"][f"{c.device} Memory"] self.assertEqual(ret["peak"], 2) - self.assertEqual(len(ret["shapes"]), 4) + self.assertEqual(len(ret["events"]), 4) if __name__ == "__main__": unittest.main() diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index c504dd191b..0354f6c733 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -154,24 +154,24 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts: events.append(struct.pack(" bytes|None: peak, mem = 0, 0 temp:dict[int, int] = {} - bufs:list[bytes] = [] - for st,_,_,e in events: + events:list[bytes] = [] + for st,_,_,e in dev_events: if not isinstance(e, ProfilePointEvent): continue if e.name == "alloc": - bufs.append(struct.pack(" peak: peak = mem if e.name == "free": - bufs.append(struct.pack(" bytes|None: # start by getting the time diffs