diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 7736eff669..0f3b75781f 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -36,6 +36,18 @@ class Mul(Function): return y*grad_output, x*grad_output register('mul', Mul) +class Div(Function): + @staticmethod + def forward(ctx, x, y): + ctx.save_for_backward(x, y) + return x / y + + @staticmethod + def backward(ctx, grad_output): + x,y = ctx.saved_tensors + return grad_output / y, -x * grad_output / y**2 +register('div', Div) + class Pow(Function): @staticmethod def forward(ctx, x, y):