diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 6ba40c0f2c..dfdf0711d1 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -124,7 +124,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u) # exclude RESHAPE/EXPAND that only serve to broadcast a CONST if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u) - if u.op in GroupOp.Movement: excluded.update(s for s in u.src if s.op is Ops.STACK and all(x.op is Ops.CONST for x in s.src)) + if u.op in {*GroupOp.Movement, Ops.PARAM}: excluded.update(s for s in u.src if s.op is Ops.STACK and all(x.op is Ops.CONST for x in s.src)) for u in toposort: argst = codecs.decode(str(u.arg), "unicode_escape") if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg)