geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

add min max rule for expand [run_process_replay] (#6434)

Browse Source
This commit is contained in:
George Hotz
2024-09-09 18:30:20 +08:00
committed by GitHub
parent 42e5c8335e
commit d5bd38c278
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/ops.py
+1
View File
@@ -417,6 +417,7 @@ class UOp(MathTrait):
if self.op is UOps.DEFINE_VAR and self.arg:
return self.arg[1].arg, self.arg[2].arg if self.arg[2].op is UOps.CONST else dtypes.max(cast(DType, self.dtype))
if self.op is UOps.RANGE: return self.src[0].vmin, (self.src[1]-1).vmax
if self.op is UOps.EXPAND: return min(x.vmin for x in self.src), max(x.vmax for x in self.src)
# TODO: UOps.SPECIAL is UOps.DEFINE_VAR
if self.op is UOps.SPECIAL: return 0, self.arg[1]-1 if isinstance(self.arg[1], int) else dtypes.max(cast(DType, self.dtype))
if self.op is UOps.CONST: return self.arg, self.arg