nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

user ceildiv in shard and fix a typo (#7690)

Browse Source
This commit is contained in:
chenyu
2024-11-13 18:25:06 -05:00
committed by GitHub
parent 51afc3cc88
commit e0ad083904
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+2 -2
View File
@@ -366,7 +366,7 @@ class Tensor(SimpleMathTrait): # pylint: disable=abstract-method
if axis < 0: axis += len(self.shape)
if splits is None:
if not isinstance(total:=self.shape[axis], int): raise RuntimeError(f"cannot shard symbolic shape {self.shape=}, {axis=}")
sz = round_up(total, len(devices)) // len(devices)
sz = ceildiv(total, len(devices))
splits = tuple([max(0, min(sz, total - sz*i)) for i in range(len(devices))])
assert sum(splits) == self.shape[axis], "specified splits do not sum up to axis shape"
boundaries = tuple(itertools.accumulate(splits))
@@ -981,7 +981,7 @@ class Tensor(SimpleMathTrait): # pylint: disable=abstract-method
```
"""
axis_arg = tuple(self._resolve_dim(x) for x in argfix(axis, *args))
if len(axis_arg) != len(dedup(axis_arg)): raise RuntimeError(f"dim can appear at least once, getting {axis_arg}")
if len(axis_arg) != len(dedup(axis_arg)): raise RuntimeError(f"dim can appear at most once, getting {axis_arg}")
return F.Flip.apply(self, axis=axis_arg)
def shrink(self, arg:Tuple[Optional[Tuple[sint, sint]], ...]) -> Tensor: