14 llops too
MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -36,10 +36,7 @@ class BinaryOps(Enum):
DIV = 3
POW = 4
A = 5
- CMP = 6
- # these are "complex" ops, and should probably be replaced
- POW_D1 = 7
- POW_D2 = 8
+ CMPEQ = 6
class ReduceOps(Enum):
SUM = 0

MODIFY
tinygrad/llops/opencl.py
tinygrad/llops/opencl.py
@@ -91,9 +91,7 @@ def binary_op(op, x, y, ret):
elif op == BinaryOps.DIV: code = "b/a"
elif op == BinaryOps.POW: code = "pow(a,b)"
elif op == BinaryOps.A: code = "a"
- elif op == BinaryOps.CMP: code = "1.0f*(a==b)"
- elif op == BinaryOps.POW_D1: code = 'b * pow(a, b-1.0f)'
- elif op == BinaryOps.POW_D2: code = 'log(a) * pow(a, b)'
+ elif op == BinaryOps.CMPEQ: code = "1.0f*(a==b)"
else: raise Exception(f"{op} isn't supported")
shape_ret, dimlist, complist = binary_broadcast(x.shape, y.shape, True)

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
@@ -62,7 +62,7 @@ class Max(Function):
def backward(ctx, grad_output):
input, ret = ctx.saved_tensors
- ret2 = binary_op(BinaryOps.CMP, input, ret, Buffer(input.shape))
+ ret2 = binary_op(BinaryOps.CMPEQ, input, ret, Buffer(input.shape))
div = reduce_op(ReduceOps.SUM, ret2, Buffer(grad_output.shape), start="1e-10")
binary_op(BinaryOps.DIV, div, ret2, ret2)
return binary_op(BinaryOps.MUL, ret2, grad_output, ret2)
@@ -112,10 +112,14 @@ class Pow(Function):
return binary_op(BinaryOps.POW, x, y, ret)
def backward(ctx, grad_output):
- x,y,ret = ctx.saved_tensors
+ x,y,powxy = ctx.saved_tensors
tmp = Buffer(grad_output.shape)
- grad_x = unbroadcast(binary_op(BinaryOps.MUL, grad_output, binary_op(BinaryOps.POW_D1, x, y, tmp), tmp), x.shape) if ctx.needs_input_grad[0] else None
- grad_y = unbroadcast(binary_op(BinaryOps.MUL, grad_output, binary_op(BinaryOps.POW_D2, x, y, tmp), tmp), y.shape) if ctx.needs_input_grad[1] else None
+ binary_op(BinaryOps.DIV, x, powxy, tmp) # pow(x,y)/x
+ binary_op(BinaryOps.MUL, y, tmp, tmp) # y * pow(x,y)/x
+ grad_x = unbroadcast(binary_op(BinaryOps.MUL, grad_output, tmp, tmp), x.shape) if ctx.needs_input_grad[0] else None
+ log_x = unary_op(UnaryOps.LOG, x, Buffer(x.shape))
+ binary_op(BinaryOps.MUL, log_x, powxy, tmp) # log(x) * pow(x,y)
+ grad_y = unbroadcast(binary_op(BinaryOps.MUL, grad_output, tmp, tmp), y.shape) if ctx.needs_input_grad[1] else None
return grad_x, grad_y
# ************* movement ops *************

14 llops too