🔨 refactor register ops (#233)

* 🔨 refactor register ops

* 🔨 reorder and register for ANE

* 🔨 refactor

* 🔨 conflicts

* 🔨 minor fix

* ane fix

* extra reshape weird

tinygrad/ops_ane.py

@@ -1,5 +1,5 @@
-from .tensor import Device, Function, register
 from functools import lru_cache
+from .tensor import Device, Function, register
 
 @lru_cache
 def compile_wrapper(ane, dat):

tinygrad/ops_cpu.py

@@ -1,5 +1,3 @@
-import sys
-import inspect
 import numpy as np
 from .tensor import Function, register
 
@@ -234,7 +232,3 @@ 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))
-
-for name, cls in inspect.getmembers(sys.modules[__name__], inspect.isclass):
-  if name[0] != "_":  register(name.lower(), cls)
-

tinygrad/ops_gpu.py

@@ -1,9 +1,7 @@
-import sys
-import inspect
+import functools
+import pyopencl as cl
 import numpy as np
 from .tensor import Function, register, GPUBuffer, Tensor, Device
-import pyopencl as cl
-import functools
 
 def buffer_new(ctx, shape, zero=False):
   return GPUBuffer(shape, hostbuf=None if not zero else np.zeros(shape, dtype=np.float32))
@@ -514,7 +512,3 @@ 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
-
-for name, cls in inspect.getmembers(sys.modules[__name__], inspect.isclass):
-  if name[0] != "_":  register(name.lower(), cls, device=Device.GPU)
-

tinygrad/tensor.py

@@ -1,9 +1,10 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
-from inspect import signature
+import sys
+import inspect
 import functools
-import numpy as np
 import os
 from collections import defaultdict
+import numpy as np
 
 # **** profiler ****
 
@@ -281,7 +282,7 @@ class Function:
   def apply(self, *x, **kwargs):
     ctx = self(*x) # self - operation i.e 'add', 'sub', etc.
     # use default params
-    params = signature(self.forward).parameters
+    params = inspect.signature(self.forward).parameters
     for p in params.values():
       if p.default is not p.empty:
         setattr(ctx, p.name, p.default)
@@ -315,11 +316,17 @@ for device in [device for device in Device.__dict__.keys() if device[0] != "_"]:
   setattr(Tensor, f"{device.lower()}_", functools.partialmethod(Tensor.to_, Device.__dict__[device]))
 
 # this registers all the operations
-import tinygrad.ops_cpu
+def _register_ops(namespace, device=Device.CPU):
+  for name, cls in inspect.getmembers(namespace, inspect.isclass):
+    if name[0] != "_":  register(name.lower(), cls, device=device)
+
+from tinygrad import ops_cpu
+_register_ops(ops_cpu)
 try:
   import pyopencl as cl
   # TODO: move this import to require_init_gpu?
-  import tinygrad.ops_gpu
+  from tinygrad import ops_gpu
+  _register_ops(ops_gpu, device=Device.GPU)
   GPU = True
 except ImportError:
   # no GPU support