From 067aeaeb2f3147a6e0d36497dcdb46009e9bc3be Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Mon, 26 Aug 2024 18:18:16 +0800
Subject: [PATCH] single arange fusion with graph rewrite (#6160)

---
 tinygrad/engine/schedule.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py
index 7ae7050a1e..d4f1f9441f 100644
--- a/tinygrad/engine/schedule.py
+++ b/tinygrad/engine/schedule.py
@@ -188,7 +188,7 @@ def push_reduceop_shape(root:UOp) -> Optional[UOp]:
 
 reduceop_fusor = PatternMatcher([
   (UPat(UOps.REDUCE_AXIS, src=(UPat(name="rsrc"), UPat(UOps.SWIZZLE, src=(UPat(name="swizzle"),))), name="root"), apply_swizzle),
-  (UPat(UOps.STORE, name="root"), push_reduceop_shape),
+  (UPat({UOps.ALU, UOps.CAST, UOps.BITCAST, UOps.STORE}, name="root"), push_reduceop_shape),
 ])
 
 def _lower_lazybuffer(outs:List[LazyBuffer], realizes:Dict[LazyBuffer, None]) -> LBScheduleItem: