Dynamically register ops to Tensor
* Dynamically register ops to Tensor
This saves lines. And reduces redundant repetition.
* ffs spacing
you don't pay me enough!
MODIFY
tinygrad/ops_cpu.py
tinygrad/ops_cpu.py
-import warnings
+import sys
+import inspect
import numpy as np
from .tensor import Function, register
@@ -14,7 +15,6 @@ class ReLU(Function):
def backward(ctx, grad_output):
input, = ctx.saved_tensors
return grad_output * (input >= 0)
-register('relu', ReLU)
class Log(Function):
@staticmethod
@@ -26,7 +26,6 @@ class Log(Function):
def backward(ctx, grad_output):
input, = ctx.saved_tensors
return grad_output / input
-register('log', Log)
class Exp(Function):
@staticmethod
@@ -39,7 +38,6 @@ class Exp(Function):
def backward(ctx, grad_output):
ret, = ctx.saved_tensors
return grad_output * ret
-register('exp', Exp)
# ************* binary ops *************
@@ -58,7 +56,6 @@ class Add(Function):
def backward(ctx, grad_output):
shape_x, shape_y = ctx.saved_tensors
return unbroadcast(grad_output, shape_x), unbroadcast(grad_output, shape_y)
-register('add', Add)
class Sub(Function):
@staticmethod
@@ -70,7 +67,6 @@ class Sub(Function):
def backward(ctx, grad_output):
shape_x, shape_y = ctx.saved_tensors
return unbroadcast(grad_output, shape_x), unbroadcast(-grad_output, shape_y)
-register('sub', Sub)
class Mul(Function):
@staticmethod
@@ -82,7 +78,6 @@ class Mul(Function):
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
return unbroadcast(y*grad_output, x.shape), unbroadcast(x*grad_output, y.shape)
-register('mul', Mul)
class Pow(Function):
@staticmethod
@@ -95,7 +90,6 @@ class Pow(Function):
x,y = ctx.saved_tensors
return unbroadcast(y * (x**(y-1.0)) * grad_output, x.shape), \
unbroadcast((x**y) * np.log(x) * grad_output, y.shape)
-register('pow', Pow)
# ************* reduce ops *************
@@ -111,7 +105,6 @@ class Sum(Function):
axis = [axis] if type(axis) is int else axis
shape = [1 if axis is None or i in axis else input.shape[i] for i in range(len(input.shape))]
return grad_output.reshape(shape) + np.zeros_like(input)
-register('sum', Sum)
class Max(Function):
@staticmethod
@@ -130,7 +123,6 @@ class Max(Function):
ret2 = (input==ret.reshape(shape))
div = ret2.sum(axis=None if axis is None else tuple(axis), keepdims=True)
return ret2*grad_output.reshape(shape)/div
-register('max', Max)
# ************* movement ops *************
@@ -151,7 +143,6 @@ class Slice(Function):
shape, = ctx.saved_tensors
narg = [(0-p[0], grad_output.shape[i]+(shape[i]-p[1])) for i,p in enumerate(ctx.arg)]
return inner_slice(grad_output, narg)
-register('slice', Slice)
class Reshape(Function):
@staticmethod
@@ -163,7 +154,6 @@ class Reshape(Function):
def backward(ctx, grad_output):
in_shape, = ctx.saved_tensors
return grad_output.reshape(in_shape)
-register('reshape', Reshape)
class Transpose(Function):
@staticmethod
@@ -174,7 +164,6 @@ class Transpose(Function):
@staticmethod
def backward(ctx, x):
return np.transpose(x, np.argsort(ctx.order))
-register('transpose', Transpose)
# ************* processing ops *************
@@ -190,7 +179,6 @@ class Matmul(Function):
grad_input = grad_output @ np.swapaxes(weight, -2, -1)
grad_weight = np.swapaxes(input, -2, -1) @ grad_output
return grad_input, grad_weight
-register('matmul', Matmul)
class Conv2D(Function):
@staticmethod
@@ -246,5 +234,6 @@ class Conv2D(Function):
gdx[:, g, :, iY:iY+H, iX:iX+W] += tg.reshape((bs, cin, H, W))
return gdx.reshape((bs, ctx.groups*cin, OY, OX)), gdw.reshape((ctx.groups*rcout, cin, H, W))
-register('conv2d', Conv2D)
+for name, cls in inspect.getmembers(sys.modules[__name__], inspect.isclass):
+ if name[0] != "_": register(name.lower(), cls)

MODIFY
tinygrad/ops_gpu.py
tinygrad/ops_gpu.py
+import sys
+import inspect
import numpy as np
from .tensor import Function, register, GPUBuffer, Tensor, Device
import pyopencl as cl
@@ -152,7 +154,6 @@ class Transpose(Function):
@staticmethod
def backward(ctx, grad_output):
return perm_axis(ctx, grad_output, np.argsort(ctx.order))
-register('transpose', Transpose, device=Device.GPU)
class Add(Function):
@staticmethod
@@ -165,7 +166,6 @@ class Add(Function):
grad_x, grad_y = grad_output, grad_output
shape_x, shape_y = ctx.saved_tensors
return unbroadcast(ctx, grad_x, shape_x), unbroadcast(ctx, grad_y, shape_y),
-register('add', Add, device=Device.GPU)
class Sub(Function):
@staticmethod
@@ -178,7 +178,6 @@ class Sub(Function):
grad_x, grad_y = grad_output, unary_op(ctx, '-a', grad_output)
shape_x, shape_y = ctx.saved_tensors
return unbroadcast(ctx, grad_x, shape_x), unbroadcast(ctx, grad_y, shape_y),
-register('sub', Sub, device=Device.GPU)
class Mul(Function):
@staticmethod
@@ -192,7 +191,6 @@ class Mul(Function):
grad_x = binary_op(ctx, 'a*b', y, grad_output)
grad_y = binary_op(ctx, 'a*b', x, grad_output)
return unbroadcast(ctx, grad_x, x.shape), unbroadcast(ctx, grad_y, y.shape),
-register('mul', Mul, device=Device.GPU)
class Pow(Function):
@staticmethod
@@ -208,7 +206,6 @@ class Pow(Function):
grad_y = binary_op(ctx, 'a*b', grad_output,
binary_op(ctx, 'pow(a, (float)b) * log(a);', x, y))
return unbroadcast(ctx, grad_x, x.shape), unbroadcast(ctx, grad_y, y.shape),
-register('pow', Pow, device=Device.GPU)
class Sum(Function):
@staticmethod
@@ -226,7 +223,6 @@ class Sum(Function):
shape = [1 if axis is None or i in axis else input.shape[i] for i in range(len(input.shape))]
output = GPUBuffer(shape, hostbuf=grad_output)
return binary_op(ctx, 'a+b', output, buffer_new(ctx, input.shape, zero=True))
-register('sum', Sum, device=Device.GPU)
class Max(Function):
@staticmethod
@@ -246,7 +242,6 @@ class Max(Function):
div = reduce_op(ctx, "out += a", "out+1e-10", ret2, axis=axis)
ret3 = binary_op(ctx, "a/b", ret2, GPUBuffer(shape, div))
return binary_op(ctx, 'a*b', ret3, GPUBuffer(shape, grad_output))
-register('max', Max, device=Device.GPU)
class Matmul(Function):
@staticmethod
@@ -301,7 +296,6 @@ class Matmul(Function):
i32(1), msize, isize, i32(1), osize, osize)
return grad_input, grad_weight
-register('matmul', Matmul, device=Device.GPU)
# ************* movement ops *************
@@ -319,7 +313,7 @@ def get_pad2d_kernel(ctx):
output[optr] = input[iptr];
}""")
-class Pad2D(Function):
+class _Pad2D(Function):
@staticmethod
def forward(ctx, x, padding=None):
bs,cin,iy,ix = x.shape
@@ -343,7 +337,6 @@ class Pad2D(Function):
i32(oy), i32(ox), i32(iy), i32(ix)
)
return ret
-#register('pad2d', Pad2D, device=Device.GPU)
class Reshape(Function):
@staticmethod
@@ -358,7 +351,6 @@ class Reshape(Function):
def backward(ctx, grad_output):
in_shape, = ctx.saved_tensors
return GPUBuffer(in_shape, hostbuf=grad_output)
-register('reshape', Reshape, device=Device.GPU)
# ************* activation ops *************
@@ -372,7 +364,6 @@ class ReLU(Function):
def backward(ctx, grad_output):
input, = ctx.saved_tensors
return binary_op(ctx, 'a * (b >= 0)', grad_output, input)
-register('relu', ReLU, device=Device.GPU)
class Log(Function):
@staticmethod
@@ -384,7 +375,6 @@ class Log(Function):
def backward(ctx, grad_output):
input, = ctx.saved_tensors
return binary_op(ctx, 'a / b', grad_output, input)
-register('log', Log, device=Device.GPU)
class Exp(Function):
@staticmethod
@@ -397,7 +387,6 @@ class Exp(Function):
def backward(ctx, grad_output):
ret, = ctx.saved_tensors
return binary_op(ctx, 'a * b', grad_output, ret)
-register('exp', Exp, device=Device.GPU)
# ************* conv ops *************
@@ -524,4 +513,6 @@ class Conv2D(Function):
convw(ctx.cl_queue, [ctx.groups*rcout*cin, H, W], None, x.cl, grad_output.cl, dw.cl, *conv_args)
convx(ctx.cl_queue, [bs, ctx.groups, cin], None, w.cl, grad_output.cl, dx.cl, *conv_args)
return dx, dw
-register('conv2d', Conv2D, device=Device.GPU)
+
+for name, cls in inspect.getmembers(sys.modules[__name__], inspect.isclass):
+ if name[0] != "_": register(name.lower(), cls, device=Device.GPU)

Dynamically register ops to Tensor (#232)

* Dynamically register ops to Tensor

This saves lines. And reduces redundant repetition.

* ffs spacing

you don't pay me enough!