scope better and don't break tests
MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
# inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
from functools import partialmethod
import numpy as np
-from numba import jit, float32
+
+# optional jit
+try:
+ from numba import jit
+except ImportError:
+ jit = lambda x: x
# **** start with two base classes ****
@@ -172,8 +177,10 @@ class LogSoftmax(Function):
register('logsoftmax', LogSoftmax)
-@jit(nopython=True)
-def conv2d_inner_forward(x, w):
+class Conv2D(Function):
+ @staticmethod
+ @jit
+ def inner_forward(x, w):
cout,cin,H,W = w.shape
ret = np.zeros((x.shape[0], cout, x.shape[2]-(H-1), x.shape[3]-(W-1)), dtype=w.dtype)
for j in range(H):
@@ -184,8 +191,9 @@ def conv2d_inner_forward(x, w):
ret[:, :, Y, X] += x[:, :, Y+j, X+i].dot(tw.T)
return ret
-@jit(nopython=True)
-def conv2d_inner_backward(grad_output, x, w):
+ @staticmethod
+ @jit
+ def inner_backward(grad_output, x, w):
dx = np.zeros_like(x)
dw = np.zeros_like(w)
cout,cin,H,W = w.shape
@@ -200,14 +208,13 @@ def conv2d_inner_backward(grad_output, x, w):
dw[:, :, j, i] += gg.T.dot(tx)
return dx, dw
-class Conv2D(Function):
@staticmethod
def forward(ctx, x, w):
ctx.save_for_backward(x, w)
- return conv2d_inner_forward(x, w)
+ return Conv2D.inner_forward(x, w)
@staticmethod
def backward(ctx, grad_output):
- return conv2d_inner_backward(grad_output, *ctx.saved_tensors)
+ return Conv2D.inner_backward(grad_output, *ctx.saved_tensors)
register('conv2d', Conv2D)

scope better and don't break tests