From dfb3e99b09d1e8ce5dcf2ae08cbaee512faf7ff4 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Fri, 25 Jul 2025 15:51:56 -0700
Subject: [PATCH] late remove ones

---
 tinygrad/opt/kernel.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tinygrad/opt/kernel.py b/tinygrad/opt/kernel.py
index 927a660bb8..fb6c7ea38d 100644
--- a/tinygrad/opt/kernel.py
+++ b/tinygrad/opt/kernel.py
@@ -451,6 +451,8 @@ class Kernel:
       ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821
       if op.op in GroupOp.Buffer and op in self.bufs:
         st = self.sts[self.bufs.index(op)]
+        # late remove all ones
+        st = st.reshape(tuple([x for x in st.shape if resolve(x != 1)]))
         # NOTE: if CONST got masked after applying opts, we create a new VALID
         if op.op is Ops.CONST and any(v.mask is not None for v in st.views): return op.view(st).valid()
         # otherwise we just replace the VIEW source