code2 was unneeded for reduce op
MODIFY
tinygrad/llops/gpu.py
tinygrad/llops/gpu.py
@@ -96,7 +96,7 @@ def binary_op(code, x, y):
prg.binop(cl_queue, [prod_list[0]] if len(dimlist) > 0 else [1], None, x.cl, y.cl, ret.cl, *dimlist, *(prod_list[1:]))
return ret
-def reduce_op(code, code2, inp, axis=None, start="0.0"):
+def reduce_op(code, inp, axis=None, start="0.0"):
if axis is None:
# full reduce
osize = [1]*len(inp.shape)
@@ -131,7 +131,7 @@ def reduce_op(code, code2, inp, axis=None, start="0.0"):
float a = a_g[idx];
"""+code+""";
}
- res_g[gid] = """+code2+""";
+ res_g[gid] = out;
}""")
reduce([np.prod(osize)], None, inp.cl,
i32(np.prod(inp.shape)//np.prod(osize)), ret.cl,

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
@@ -33,7 +33,7 @@ class Exp(UnaryOp):
class Sum(Function):
def forward(ctx, input, axis=None):
ctx.save_for_backward(input.shape)
- return reduce_op("out += a", "out", input, axis=axis)
+ return reduce_op("out += a", input, axis=axis)
def backward(ctx, grad_output):
shape_input, = ctx.saved_tensors
@@ -42,14 +42,14 @@ class Sum(Function):
class Max(Function):
def forward(ctx, input, axis=None):
- ret = reduce_op("out = max(a,out)", "out", input, axis=axis, start="-INFINITY")
+ ret = reduce_op("out = max(a,out)", input, axis=axis, start="-INFINITY")
ctx.save_for_backward(input, axis, ret)
return ret
def backward(ctx, grad_output):
input, axis, ret = ctx.saved_tensors
ret2 = binary_op("1.0*(a==b)", input, ret)
- div = reduce_op("out += a", "out+1e-10", ret2, axis=axis)
+ div = reduce_op("out += a", ret2, axis=axis, start="1e-10")
ret3 = binary_op("a/b", ret2, div)
return binary_op('a*b', ret3, grad_output)
@@ -57,7 +57,7 @@ class Max(Function):
def unbroadcast(out, in_sh):
sum_axis = [i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1] if in_sh != (1,) else None
- return reduce_op("out += a", "out", out, sum_axis)
+ return reduce_op("out += a", out, sum_axis)
class Add(Function):
def forward(ctx, x, y):

code2 was unneeded for reduce op