github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

clean up reduce a little

Browse Source
This commit is contained in:
George Hotz
2022-06-09 09:34:19 -07:00
parent d3e7238bdd
commit 259b536e3a
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tinygrad/llops/ops_cpu.py
View File

@@ -37,14 +37,17 @@ def binary_op(op, x, y, ret):
return ret
def reduce_op(op, inp, ret):
if inp.shape == ret.shape: # this is just a copy
if inp.shape == ret.shape: # this is just a copy, regardless of the reduce op
ret[:] = inp
return ret
if ret.shape == (1,): axis=tuple(range(len(inp.shape)))
else: axis = tuple([i for i,(a,b) in enumerate(zip(inp.shape, ret.shape)) if a != b])
if op == ReduceOps.SUM: ret[:] = inp.sum(axis, keepdims=True)
elif op == ReduceOps.MAX: ret[:] = inp.amax(axis, keepdims=True)
else: raise Exception(f"{op} isn't supported")
else:
if ret.shape == (1,): # full reduce
axis = tuple(range(len(inp.shape)))
else:
assert len(inp.shape) == len(ret.shape)
axis = tuple([i for i,(a,b) in enumerate(zip(inp.shape, ret.shape)) if a != b])
if op == ReduceOps.SUM: ret[:] = inp.sum(axis, keepdims=True)
elif op == ReduceOps.MAX: ret[:] = inp.amax(axis, keepdims=True)
else: raise Exception(f"{op} isn't supported")
return ret
def reshape(x, shape):