From dd16d5aead62e0207c0c3c50c19bc8b67e176c55 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Tue, 28 Jul 2026 19:03:06 -0700
Subject: [PATCH] apply shrink bugfix for 3.11 (#17271)

---
 .github/workflows/test.yml |  2 ++
 tinygrad/schedule/multi.py | 13 +++++++------
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 4fa90bb270..a7f37e8250 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -167,6 +167,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         deps: testing_unit
         llvm: 'true'
         amd: 'true'
@@ -200,6 +201,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         pydeps: "pre-commit"
         deps: testing_unit
         llvm: 'true'
diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py
index c9da55217a..ab01eb3778 100644
--- a/tinygrad/schedule/multi.py
+++ b/tinygrad/schedule/multi.py
@@ -5,17 +5,18 @@ from tinygrad.schedule.allreduce import handle_allreduce
 
 # ***** multi rewrite MSELECT/MSTACK *****
 
+def _apply_shrink(marg, s:UOp, i:int) -> UOp:
+  new_arg = [tuple([x.substitute({drng[0]:drng[0].const_like(i)}) if isinstance(x, UOp) and
+                    (drng:=[r for r in x.ranges if r.arg[-1] is AxisType.DEVICE]) else x for x in ss]) for ss in marg]
+  return s._mop(Ops.SHRINK, tuple(new_arg))
+
 def mstack_early_shrink(ms:UOp, shrink:UOp):
   ret:list[UOp] = []
-  def apply_shrink(s:UOp, i:int) -> UOp:
-    new_arg = [tuple([x.substitute({drng[0]:drng[0].const_like(i)}) if isinstance(x, UOp) and
-                      (drng:=[r for r in x.ranges if r.arg[-1] is AxisType.DEVICE]) else x for x in ss]) for ss in shrink.marg]
-    return s._mop(Ops.SHRINK, tuple(new_arg))
   for i, x in enumerate(ms.src):
     if x.op is Ops.COPY:
-      ret.append(apply_shrink(x.src[0], i).copy_to_device(x.device))
+      ret.append(_apply_shrink(shrink.marg, x.src[0], i).copy_to_device(x.device))
     else:
-      ret.append(apply_shrink(x, i).contiguous())
+      ret.append(_apply_shrink(shrink.marg, x, i).contiguous())
   return ms.replace(src=tuple(ret))
 
 def lower_broadcast_copy(c:UOp, x:UOp):