diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 1d3d1f5b78..09a1dbe497 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -2188,7 +2188,7 @@ class Tensor(OpMixin): assert self.shape[axis] != 0 and op in (Ops.ADD, Ops.MAX, Ops.MUL) pl_sz = self.shape[axis] - int(not _include_initial) pooled = self.transpose(axis,-1).pad((pl_sz, -int(_include_initial)), value=identity_element(op, self.dtype))._pool((self.shape[axis],)) - return {Ops.ADD: pooled.sum(-1), Ops.MAX: pooled.max(-1), Ops.MUL: pooled.prod(-1)}[op].transpose(axis, -1) + return getattr(pooled, {Ops.ADD: "sum", Ops.MAX: "max", Ops.MUL: "prod"}[op])(-1).transpose(axis, -1) def _split_cumalu(self, axis:int, op:Ops) -> Tensor: axis = self._resolve_dim(axis) @@ -2201,8 +2201,7 @@ class Tensor(OpMixin): base = ret[..., -1]._cumalu(-1, op, _include_initial=True) base = base.unsqueeze(-1).expand(*base.shape, ret.shape[-1]) def fix(x: Tensor) -> Tensor: return x.flatten(start_dim=-2)[..., -s:].transpose(axis,-1) - reduce_fxns: dict[Ops, Callable[[Tensor, Tensor], Tensor]] = {Ops.ADD: Tensor.__add__, Ops.MAX: Tensor.maximum, Ops.MUL: Tensor.__mul__} - return reduce_fxns[op](fix(ret), fix(base)) + return getattr(fix(ret), {Ops.ADD: "add", Ops.MAX: "maximum", Ops.MUL: "mul"}[op])(fix(base)) def cumsum(self, axis:int=0) -> Tensor: """