github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-19 02:44:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add cumsum with n-dim inputs (#922)

Browse Source 
* add cumsum with n-dim inputs, over arbitrary axis + relevant tests

* increased rtol for cumsum test

* move test_cumsum into test_ops

* skip arange test for images as relies on cumsum

* Fix typo

* rewrite cumsum to work with images
This commit is contained in:
Tom Edwards
2023-06-05 00:55:23 +01:00
committed by GitHub
parent 091bd65a68
commit 5bbcbd145c
2 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tinygrad/tensor.py
View File

@@ -460,9 +460,10 @@ class Tensor:
r = (x*w).sum(-1)
return r.reshape((*r.shape[:-2], r.shape[-1])) if len(self.shape) == 1 else r
# TODO: make this work for n-dimensional inputs
def cumsum(self): return self.reshape(1, 1, 1, self.shape[0]).conv2d(Tensor.ones(1, 1, 1, self.shape[0]), padding=(self.shape[0] - 1, 0, 0, 0)).flatten()
def cumsum(self, axis=0):
x = self.permute(*(i for i in range(self.ndim) if i != axis), axis)
return x.reshape(1, 1, -1, self.shape[axis]).conv2d(Tensor.ones(1, 1, 1, self.shape[axis]), padding=(self.shape[axis]-1, 0, 0, 0)).reshape(*x.shape).permute(*range(axis), self.ndim - 1, *range(axis, self.ndim-1))
# ***** mlops (unary) *****
def contiguous(self): return mlops.Contiguous.apply(self)