tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

simpler _split_cumalu [PR] (#17180)

Browse Source
This commit is contained in:
chenyu
2026-07-24 10:52:50 -04:00
committed by GitHub
parent fcb1d3fc9f
commit e0b30b7a1d
1 changed files with 3 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/mixin/op.py
+3 -5
View File
@@ -738,11 +738,9 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
SPLIT = 256
value = identity_element(op, self.dtype)
if not isinstance(s:=self.shape[axis], int) or s <= SPLIT*2: return self._cumalu(axis, op)
ret = self.transpose(axis,-1)._pad_constant((None,)*(self.ndim-1)+((round_up(s,SPLIT)-s,0),), value).unflatten(-1,(-1,SPLIT))._cumalu(-1, op)
base = ret[..., -1]._cumalu(-1, op)._pad_constant((None,)*(ret.ndim-2) + ((1, -1),), value)
base = base.unsqueeze(-1).expand(*base.shape, ret.shape[-1])
def fix(x: Self) -> Self: return x.flatten(start_dim=-2)[..., -s:].transpose(axis,-1)
return fix(ret).alu(op, fix(base))
chunks = self.transpose(axis,-1)._pad_constant((None,)*(self.ndim-1)+((round_up(s,SPLIT)-s,0),), value).unflatten(-1,(-1,SPLIT))._cumalu(-1, op)
base = chunks[..., -1]._cumalu(-1, op)._pad_constant((None,)*(chunks.ndim-2) + ((1, -1),), value)
return chunks.alu(op, base.unsqueeze(-1)).flatten(start_dim=-2)[..., -s:].transpose(axis,-1)
def cumsum(self, axis:int=0) -> Self:
"""