From ece07a3d128519633186ca76b0aabbfb0df4dbcc Mon Sep 17 00:00:00 2001
From: NeuralLink <38289341+kartik4949@users.noreply.github.com>
Date: Sat, 2 Jan 2021 21:17:16 +0530
Subject: [PATCH] :hammer: refactor register ops (#233)

* :hammer: refactor register ops

* :hammer: reorder and register for ANE

* :hammer: refactor

* :hammer: conflicts

* :hammer: minor fix

* ane fix

* extra reshape weird
---
 tinygrad/ops_ane.py |  2 +-
 tinygrad/ops_cpu.py |  6 ------
 tinygrad/ops_gpu.py | 10 ++--------
 tinygrad/tensor.py  | 17 ++++++++++++-----
 4 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/tinygrad/ops_ane.py b/tinygrad/ops_ane.py
index 434fd13190..c95dbdc413 100644
--- a/tinygrad/ops_ane.py
+++ b/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):
diff --git a/tinygrad/ops_cpu.py b/tinygrad/ops_cpu.py
index 8db8ceaec8..879d225a5e 100644
--- a/tinygrad/ops_cpu.py
+++ b/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)
-
diff --git a/tinygrad/ops_gpu.py b/tinygrad/ops_gpu.py
index 255127199d..2f757c8faa 100644
--- a/tinygrad/ops_gpu.py
+++ b/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)
-
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 60b94dabee..03de64b37c 100644
--- a/tinygrad/tensor.py
+++ b/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