diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index a17a5f42ac..a305cd3fc4 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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: