From d5bd38c278eeebc6bc4ecfb039a3a98e9a2bbc49 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Mon, 9 Sep 2024 18:30:20 +0800
Subject: [PATCH] add min max rule for expand [run_process_replay] (#6434)

---
 tinygrad/ops.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index 7b5cd70bfa..eedde4cabf 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -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