sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

Tensor cumalu cleanups (#15579)

Browse Source 
* Tensor cumalu cleanups

* happy
This commit is contained in:
chenyu
2026-04-02 05:23:22 -04:00
committed by GitHub
parent 1aa04eab08
commit 61bc91aa8c
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+2 -3
View File
@@ -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:
"""