reduce axis at the end
MODIFY
test/graph_batchnorm.py
test/graph_batchnorm.py
from tinygrad.tensor import Tensor
-from tinygrad.nn import Conv2d, BatchNorm2D
-from tinygrad import optim
+from tinygrad.nn import Conv2d, BatchNorm2D, optim
from extra.utils import get_parameters # TODO: move to optim
import unittest

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -114,8 +114,6 @@ class GPUBuffer:
def movement_op(x, op:MovementOps, arg) -> GPUBuffer: return type(x)(ShapeTracker(x.st).movement_op(op, arg), x)
def reduce_op(x, op:ReduceOps, new_shape:Tuple[int, ...]): return type(x)(new_shape)._processing_op([("A", x)], code="acc", earlycode=GPUBuffer.code_for_op[op], earlybufs=set("A"), start=GPUBuffer.start_for_op[op])
- #REQUIRES_SIMPLE_REDUCE = True
- # is there a downside to REQUIRES_SIMPLE_REDUCE always?
def _processing_op(ret, bufs: List[Tuple[str, GPUBuffer]]=[], code:str="acc", C:Optional[ConvArgs]=None, start="0.0", reduce_shape=None, earlybufs:Set[str]=set(), earlycode:str="acc") -> GPUBuffer:
assert C is None

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -247,13 +247,11 @@ class LazyBuffer:
# TODO: permute to put all the reduce axis at the end
def reduce_op(self:LazyBuffer, op:ReduceOps, new_shape:Tuple[int, ...]) -> LazyBuffer:
if self.shape == tuple(new_shape): return self
- if getattr(self.dbuffer, "REQUIRES_SIMPLE_REDUCE", False) and (len(new_shape) != 2 or new_shape[1] != 1):
- num, red = prod([s for s,n in zip(self.shape, new_shape) if n != 1]), prod([s for s,n in zip(self.shape, new_shape) if n == 1])
+ # move the reduce axes to the end
x = self.movement_op(MovementOps.PERMUTE, [i for i,n in enumerate(new_shape) if n != 1] + [i for i,n in enumerate(new_shape) if n == 1])
- x = x.movement_op(MovementOps.RESHAPE, (num, red)) # remove this reshape, at the end is enough
- return x.reduce_op(op, (num, 1)).movement_op(MovementOps.RESHAPE, new_shape)
- else:
- return LazyBuffer(self.device, tuple(new_shape), ReduceOps, LazyOp(op, (self,), tuple(new_shape)))
+ new_tmp_shape = tuple([n for n in new_shape if n != 1] + [1 for n in new_shape if n == 1])
+ # NOTE: this reshape can only move around 1s
+ return LazyBuffer(x.device, new_tmp_shape, ReduceOps, LazyOp(op, (x,), new_tmp_shape)).movement_op(MovementOps.RESHAPE, new_shape)
# syntactic sugar around PAD and SHRINK
# TODO: turn RESHAPE into EXPAND and CONTRACT (current EXPAND should be REPEAT)

reduce axis at the end