pow+div on GPU
MODIFY
tinygrad/opsgpu.py
tinygrad/opsgpu.py
@@ -85,6 +85,22 @@ class Mul(Function):
binary_op(ctx, 'res_g[gid] = a_g[gid] * b_g[gid];', x, grad_output)
register('mul', Mul, gpu=True)
+class Pow(Function):
+ @staticmethod
+ def forward(ctx, x, y):
+ ctx.save_for_backward(x, y)
+ return binary_op(ctx, 'res_g[gid] = pow(a_g[gid], b_g[gid]);', x, y)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ x,y = ctx.saved_tensors
+ gradx = binary_op(ctx, 'res_g[gid] = a_g[gid] * b_g[gid];', grad_output,
+ binary_op(ctx, 'res_g[gid] = b_g[gid] * (pow((float)a_g[gid], (float)(b_g[gid]-1.0)));', x, y))
+ grady = binary_op(ctx, 'res_g[gid] = a_g[gid] * b_g[gid];', grad_output,
+ binary_op(ctx, 'res_g[gid] = pow((float)a_g[gid], (float)b_g[gid]) * log(a_g[gid]);', x, y))
+ return gradx, grady
+register('pow', Pow, gpu=True)
+
class Sum(Function):
@staticmethod
def forward(ctx, input):

pow+div on GPU (#57)