geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

E501

Browse Source
This commit is contained in:
David Hou
2024-02-21 23:33:35 -08:00
parent fb3db11b8e
commit 5e8b6f8d10
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/features/multi.py
+1 -1
View File
@@ -101,7 +101,7 @@ class MultiLazyBuffer:
# new_axis is the last one that preserves prod(prior to new_axis) and must not move items between shards
new_axis = len(arg_acc) - arg_acc[::-1].index(prod(self.shape[:self.axis])) - 1
if not arg[new_axis] == self.shape[self.axis]:
assert self.shape[self.axis] % len(self.real_lbs) == 0, f"cannot reshape on-axis for uneven shards {self.axis} {self.shape} {len(self.real_lbs)}"
assert self.shape[self.axis] % len(self.real_lbs) == 0, f"cannot reshape on-axis for uneven shard {self.axis} {self.shape} {len(self.real_lbs)}"
assert arg[new_axis] % len(self.real_lbs) == 0, f"new on-axis shape must divide evenly between devices {new_axis} {arg} {len(self.real_lbs)}"
return MultiLazyBuffer([x.reshape(tuple(s if a != new_axis else
x.shape[self.axis] if s == self.shape[self.axis] else