diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 91a2110122..40122499c9 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -91,8 +91,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @functools.cached_property def key(self) -> bytes: return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest() - def __repr__(self): return pretty_print(self, lambda x: f"{type(self).__name__}({x.op}, {x.dtype}, arg={x.argstr()}, src=(%s))") + def __repr__(self): return pretty_print(self, lambda x: f"{type(self).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=(%s))") def argstr(self): return f'({", ".join(map(str, self.arg))})' if self.op is Ops.REDUCE_AXIS else repr(self.arg) + def tagstr(self): return f", tag={self.tag}" if self.tag is not None else "" @functools.cached_property def parents(self:UOp) -> dict[UOp, None]: