mlops cleanup
MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -5,38 +5,32 @@ from tinygrad.tensor import Function
# ************* unary ops *************
-class _UnaryOp(Function):
+class ReLU(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
- return input.unary_op(ctx.fop)
+ return input.unary_op(UnaryOps.RELU)
def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- return input.binary_op(ctx.bop, grad_output)
-
-class ReLU(_UnaryOp):
- fop = UnaryOps.RELU
-
- def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- ret = input.unary_op(UnaryOps.SIGN)
+ ret = ctx.saved_tensors[0].unary_op(UnaryOps.SIGN)
ret = ret.unary_op(UnaryOps.RELU)
return ret.binary_op(BinaryOps.MUL, grad_output)
-class Log(_UnaryOp):
- fop = UnaryOps.LOG
+class Log(Function):
+ def forward(ctx, input):
+ ctx.save_for_backward(input)
+ return input.unary_op(UnaryOps.LOG)
def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- return grad_output.binary_op(BinaryOps.DIV, input)
+ return grad_output.binary_op(BinaryOps.DIV, ctx.saved_tensors[0])
-class Exp(_UnaryOp):
+class Exp(Function):
def forward(ctx, input):
ret = input.unary_op(UnaryOps.EXP)
ctx.save_for_backward(ret) # we save the output here, not the input
return ret
- bop = BinaryOps.MUL
+ def backward(ctx, grad_output):
+ return ctx.saved_tensors[0].binary_op(BinaryOps.MUL, grad_output)
# TODO: add Neg? confirm the optimizer on Sub good enough
@@ -44,12 +38,11 @@ class Exp(_UnaryOp):
class Sum(Function):
def forward(ctx, input, axis=None):
- ctx.save_for_backward(input.shape)
+ ctx.input_shape = input.shape
return input.reduce_op(ReduceOps.SUM, reduce_shape(input.shape, axis))
def backward(ctx, grad_output):
- shape_input, = ctx.saved_tensors
- return grad_output.movement_op(MovementOps.EXPAND, shape_input)
+ return grad_output.movement_op(MovementOps.EXPAND, ctx.input_shape)
class Max(Function):
def forward(ctx, input, axis=None):
@@ -95,9 +88,8 @@ class Mul(Function):
return x.binary_op(BinaryOps.MUL, y)
def backward(ctx, grad_output):
- x,y = ctx.saved_tensors
- grad_x = y.binary_op(BinaryOps.MUL, grad_output) if ctx.needs_input_grad[0] else None
- grad_y = x.binary_op(BinaryOps.MUL, grad_output) if ctx.needs_input_grad[1] else None
+ grad_x = ctx.saved_tensors[1].binary_op(BinaryOps.MUL, grad_output) if ctx.needs_input_grad[0] else None
+ grad_y = ctx.saved_tensors[0].binary_op(BinaryOps.MUL, grad_output) if ctx.needs_input_grad[1] else None
return grad_x, grad_y
# TODO: add Div? is the optimizer on Pow good enough?
@@ -125,53 +117,46 @@ class Pow(Function):
# NOTE: this is sum in reverse
class Expand(Function):
def forward(ctx, x, shape):
- ctx.save_for_backward(x.shape)
+ ctx.input_shape = x.shape
return x.movement_op(MovementOps.EXPAND, shape)
def backward(ctx, grad_output):
- in_shape, = ctx.saved_tensors
- return grad_output.reduce_op(ReduceOps.SUM, in_shape)
+ return grad_output.reduce_op(ReduceOps.SUM, ctx.input_shape)
class Reshape(Function):
def forward(ctx, x, shape):
- ctx.save_for_backward(x.shape)
+ ctx.input_shape = x.shape
shape = tuple(-prod(x.shape) // prod(shape) if s == -1 else s for s in shape)
return x.movement_op(MovementOps.RESHAPE, shape)
def backward(ctx, grad_output):
- in_shape, = ctx.saved_tensors
- return grad_output.movement_op(MovementOps.RESHAPE, in_shape)
+ return grad_output.movement_op(MovementOps.RESHAPE, ctx.input_shape)
class Permute(Function):
def forward(ctx, x, order=(1,0)):
- ctx.save_for_backward(order)
+ ctx.input_order = order
return x.movement_op(MovementOps.PERMUTE, order)
def backward(ctx, grad_output):
- order, = ctx.saved_tensors
- norder = np.argsort(order).tolist()
+ norder = np.argsort(ctx.input_order).tolist()
return grad_output.movement_op(MovementOps.PERMUTE, norder)
# TODO: merge Slice and Flip into Stride with the 3 arguments
-
class Slice(Function):
def forward(ctx, x, arg=None):
- ctx.save_for_backward(x.shape, arg)
+ ctx.narg = [(0-p[0], x.shape[i]-p[0]) for i,p in enumerate(arg)]
return x.movement_op(MovementOps.SLICE, arg)
def backward(ctx, grad_output):
- shape, arg = ctx.saved_tensors
- narg = [(0-p[0], grad_output.shape[i]+(shape[i]-p[1])) for i,p in enumerate(arg)]
- return grad_output.movement_op(MovementOps.SLICE, narg)
+ return grad_output.movement_op(MovementOps.SLICE, ctx.narg)
class Flip(Function):
def forward(ctx, x, axis):
- ctx.save_for_backward(axis)
+ ctx.axis = axis
return x.movement_op(MovementOps.FLIP, axis)
def backward(ctx, grad_output):
- axis, = ctx.saved_tensors
- return grad_output.movement_op(MovementOps.FLIP, axis)
+ return grad_output.movement_op(MovementOps.FLIP, ctx.axis)
# ************* processing ops *************
@@ -182,12 +167,13 @@ class Conv2D(Function):
return x.processing_op(ProcessingOps.CONV, w, C)
def forward(ctx, x, w, stride=1, groups=1, dilation=1, padding=0):
- C = get_conv_args(x.shape, w.shape, stride, groups, dilation=dilation, padding=padding)
- ctx.save_for_backward(x,w,C)
- return ctx._conv(x, w, C)
+ ctx.C = get_conv_args(x.shape, w.shape, stride, groups, dilation=dilation, padding=padding)
+ ctx.save_for_backward(x,w)
+ return ctx._conv(x, w, ctx.C)
def backward(ctx, grad_output):
- x, w, C = ctx.saved_tensors
+ x, w = ctx.saved_tensors
+ C = ctx.C # conv args from the context
dx, dw = None, None
if ctx.needs_input_grad[0]: # compute derivative of inputs using ProcessingOps.CONV (this is a transposed conv)
xt = grad_output

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -183,6 +183,7 @@ class Tensor:
return self[:, :, -padding[2]:self.shape[2]+padding[3], -padding[0]:self.shape[3]+padding[1]]
def matmul(x, w):
+ # NOTE: we use a 1x1 conv2d to do the matmul. mxk @ kxn = (1,k,m,1).conv2d(n,k,1,1)
bs, groups = prod(x.shape[0:-2]), prod(w.shape[0:-2])
cin, cout = w.shape[-2], w.shape[-1]
out_shape_t = tuple(list(x.shape[0:-2])+[cout,-1])
@@ -321,9 +322,7 @@ class Tensor:
ret = self.mul(weight.reshape(shape=shp)) if len(weight.shape) == 1 else self.dot(weight)
return ret.add(bias.reshape(shape=shp))
- def sequential(self, ll):
- for l in ll: self = l(self)
- return self
+ def sequential(self, ll): return functools.reduce(lambda x,f: f(x), ll, self)
def layernorm(x, eps=1e-5):
y = (x - x.mean(axis=-1, keepdim=True))

mlops cleanup