print tag in UOp [pr] (#10755)

This commit is contained in:
qazal
2025-06-10 21:16:07 +03:00
committed by GitHub
parent 14fa62c61d
commit 9e1d1ebc52
+2 -1
View File
@@ -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]: