From 61de654efad3ea03a830d735bb7093218a5a13ff Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Tue, 4 Feb 2025 13:22:31 -0500
Subject: [PATCH] minor shard cleanup [pr] (#8888)

---
 tinygrad/ops.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index 896dae1d1d..dd2978b628 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -466,12 +466,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if axis is None: lbs = [self] * len(devices)
     else:
       if self.shape[axis] % len(devices) != 0: raise RuntimeError(f"multi axis uneven: {self.shape[axis]=} {axis=} {len(devices)=}")
+      # NOTE: this works for both even shards and uneven shards
       sz = self.shape[axis] // len(devices)
       sizes = [max(0, min(sz, self.shape[axis] - sz*i)) for i in range(len(devices))]
-      lbs, off = [], 0
-      for sz in sizes:
+      lbs = []
+      for sz,off in zip(sizes, itertools.accumulate(sizes, initial=0)):
         lbs.append(self.shrink(tuple((0,s) if i != axis else (off,off+sz) for i,s in enumerate(self.shape))))
-        off += sz
     sharded_lbs = [lb.copy_to_device(d) for lb,d in zip(lbs, devices)]
     # NOTE: this contiguous is making it impossible for the scheduler to do late const folding
     return UOp.multi(*[lb.contiguous() for lb in sharded_lbs], axis=axis)