cpu line savings and cleaner
MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
@@ -7,7 +7,8 @@ class CPUBuffer(np.ndarray):
UnaryOps.NOOP: lambda x: x[:], UnaryOps.NEG: lambda x: -x, UnaryOps.RELU: lambda x: x.relu(),
UnaryOps.EXP: lambda x: x.exp(), UnaryOps.LOG: lambda x: x.log(), UnaryOps.SIGN: lambda x: x.sign(), UnaryOps.RECIPROCAL: lambda x: 1.0/x,
BinaryOps.ADD: operator.add, BinaryOps.SUB: operator.sub, BinaryOps.MUL: operator.mul,
- BinaryOps.DIV: operator.truediv, BinaryOps.POW: operator.pow, BinaryOps.CMPEQ: lambda x,y: (x==y).float()
+ BinaryOps.DIV: operator.truediv, BinaryOps.POW: operator.pow, BinaryOps.CMPEQ: lambda x,y: (x==y).float(),
+ ReduceOps.SUM: lambda x, axis: x.sum(axis, keepdims=True), ReduceOps.MAX: lambda x, axis: x.amax(axis, keepdims=True)
}
def relu(x): return np.maximum(x, 0)
@@ -18,7 +19,7 @@ class CPUBuffer(np.ndarray):
def flip(x, axis): return np.flip(x, axis)
def amax(x, *args, **kwargs): return np.amax(x, *args, **kwargs)
def permute(x, order): return x.transpose(order)
- def pad(x, padding): return np.pad(x, padding).view(CPUBuffer) if any(x != 0 or y != 0 for x,y in padding) else x
+ def pad(x, padding): return np.pad(x, padding).view(CPUBuffer)
def expand(x, new_shape): return np.broadcast_to(x, new_shape).view(CPUBuffer)
def as_strided(x, size, stride): return np.lib.stride_tricks.as_strided(x, shape=size, strides=[y*x.dtype.itemsize for y in stride]).view(CPUBuffer)
def contiguous(x): return x.ravel().reshape(x.shape)
@@ -33,12 +34,7 @@ class CPUBuffer(np.ndarray):
def reduce_op(x, op, new_shape):
assert len(x.shape) == len(new_shape)
axis = tuple([i for i,(a,b) in enumerate(zip(x.shape, new_shape)) if a != b])
- if x.shape == new_shape:
- return x[:] # this is just a copy, regardless of the reduce op
- elif op == ReduceOps.SUM:
- return x.sum(axis, keepdims=True)
- elif op == ReduceOps.MAX:
- return x.amax(axis, keepdims=True)
+ return CPUBuffer.fxn_for_op[op](x, axis) if len(axis) > 0 else x[:]
def movement_op(x, op, arg=None):
if op == MovementOps.SHRINK:

cpu line savings and cleaner