diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index b3106c0241..7ac23c89f8 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -79,10 +79,10 @@ def uop_to_json(x:UOp) -> dict[int, dict]: arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(x.dtype) else f"{x.arg}" label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "") try: + if len(rngs:=u.ranges): + label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})" if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: label += f"\n{shape_to_str(u.shape)}" - elif len(rngs:=u.ranges): - label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})" if u.op is Ops.INDEX: label += f"\n{u.render()}" except Exception: