Mypy fun (#541)

* mypy fun

* things are just faster

* running fast

* mypy is fast

* compile.sh

* no gpu hack

* refactor ops_cpu and ops_torch to not subclass

* make weak buffer work

* tensor works

* fix test failing

* cpu/torch cleanups

* no or operator on dict in python 3.8

* that was junk

* fix warnings

* comment and touchup

.gitignore

@@ -3,6 +3,7 @@ notebooks
 .*.swp
 .*.swo
 *.pyc
+*.so
 build
 dist
 *.egg-info

compile.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+mypyc tinygrad/llops/ops_gpu.py tinygrad/shape/__init__.py tinygrad/ops.py tinygrad/ast.py \
+  tinygrad/helpers.py tinygrad/mlops.py tinygrad/nn/__init__.py tinygrad/graph.py tinygrad/lazy.py \
+  tinygrad/tensor.py 
+

examples/benchmark_train_efficientnet.py

@@ -66,7 +66,7 @@ if __name__ == "__main__":
       CL.CACHE = None
 
     mem_used = CL.mem_used
-    loss_cpu = loss.detach().cpu().data[0]
+    loss_cpu = loss.detach().numpy()[0]
     cl = time.monotonic()
 
     print(f"{(st-cpy)*1000.0:7.2f} ms cpy,  {(cl-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {(et-mt)*1000.0:7.2f} ms realize, {(cl-et)*1000.0:7.2f} ms CL, {loss_cpu:7.2f} loss, {tensors_allocated():4d} tensors, {mem_used/1e9:.2f} GB used, {ops*1e-9/(cl-st):9.2f} GFLOPS")

extra/lib_test_ast.py

@@ -3,8 +3,7 @@ import numpy as np
 from typing import Dict, Type
 from tinygrad.ast import ASTKernel
 from tinygrad.llops.ops_cpu import CPUBuffer
-from tinygrad.ops import DeviceBuffer
-from tinygrad.lazy import map_buffers
+from tinygrad.ops import DeviceBuffer, map_buffers
 
 in_test = False
 test_cnt = 0

rmso.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+rm tinygrad/*.so tinygrad/shape/*.so tinygrad/llops/*.so tinygrad/nn/*.so tinygrad/runtime/*.so *.so
+

tinygrad/ast.py

@@ -1,6 +1,6 @@
-from enum import Enum
+from enum import Enum, auto
 import itertools
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 from tinygrad.helpers import prod, dedup, all_same
 from tinygrad.ops import LazyOp, MovementOps, get_lazyop_info, get_buffers, ReduceOps, get_lazyops
 from tinygrad.shape import ShapeTracker, View, strides_for_shape
@@ -11,14 +11,15 @@ def get_first_reduce(shapes):
       return i
   return len(shapes[0])  # off the end
 
-Types = Enum("Types", ["FLOAT", "FLOAT4"])
+# this will be removed soon anyway
+class Types(Enum): FLOAT = auto(); FLOAT4 = auto() # noqa: E702
 class Token:
   def __init__(self, tok:str, typ:Types, ptr:bool=False):
     assert isinstance(tok, str)
     self.tok, self.typ, self.ptr = tok, typ, ptr
     self.axis : List[Tuple[int, int, bool]] = []
   def array(self, length, stride, reduce): self.axis.append((length, stride, reduce))
-  def size(self): return prod(x[0] for x in self.axis)
+  def size(self): return prod([x[0] for x in self.axis])
   def offsets(self): return [sum(t) for t in itertools.product(*[[y*x[1] for y in range(x[0])] for x in self.axis[::-1]])] if len(self.axis) else [0]
   # TODO: this is sort of a hack, it gets the accumulator indices
   def acc_offsets(self):
@@ -53,12 +54,13 @@ class ASTKernel:
     # TODO: should be optional if it's hitting a function cache
     self.processed = False
 
-  def process(self):
+  def process(self) -> None:
     if self.processed: return
     self.processed = True
     reduceops = [x for x in get_lazyops(self.ast) if x.op in ReduceOps]
     assert len(dedup(reduceops)) <= 1, "max one reduce op in an ast"
     self.reduceop = reduceops[0] if reduceops else None
+    self.reduceopop : Optional[ReduceOps] = self.reduceop.op if self.reduceop is not None and isinstance(self.reduceop.op, ReduceOps) else None
     self.earlybufs = dedup(get_buffers(self.reduceop)) if self.reduceop else []
 
     self.buftokens = [Token(f"data{i}", Types.FLOAT, ptr=True) for i in range(len(self.bufs))]
@@ -98,7 +100,7 @@ class ASTKernel:
     print_ast(self.input_ast, "ast")
 
   @property
-  def shape_len(self): return len(self.sts[0].shape)
+  def shape_len(self) -> int: return len(self.sts[0].shape)
 
   def simplify_ones(self):
     # remove places where the shape is all ones

tinygrad/lazy.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
-from typing import Optional, Tuple, Union, List, Dict, Any, Final
+from typing import Optional, Tuple, Union, List, Dict, Any, ClassVar
 import sys, weakref
+from weakref import WeakValueDictionary
 from tinygrad.helpers import ConvArgs, prod
 from tinygrad.shape import ShapeTracker
-from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps, OpType, LazyOp, get_buffers, DEBUG
+from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps, OpType, LazyOp, get_buffers, map_buffers, DEBUG
 from tinygrad.graph import log_op
 from tinygrad.helpers import getenv
 
@@ -22,12 +23,6 @@ REMOVE_MOVEMENT_NOPS, MERGE_UNARY_OPS, MERGE_ELEMENTWISE_INTO_REDUCE, SHUFFLE_MO
 MERGE_ELEMENTWISE_OPS, MERGE_ONE_REDUCE_INTO_ELEMENTWISE = OPT>=2, OPT>=2
 SHUFFLE_PAD_OPS = OPT>=3  # NOTE: 0/0 is NaN if you pad, so this can change the output
 
-# **** realize helpers ****
-def map_buffers(real_srcs, x:LazyOp) -> LazyOp:
-  if x in real_srcs:
-    return map_buffers(real_srcs, real_srcs[x]) if isinstance(real_srcs[x], LazyOp) else real_srcs[x]
-  return LazyOp(x.op, tuple(map_buffers(real_srcs, y) for y in x.src), x.arg)
-
 # **** realize functions ****
 def _ast_reduceops(self:LazyBuffer) -> LazyOp:
   # TODO: this can also corealize a binary op after the reduce, not just before
@@ -74,18 +69,25 @@ def get_weakop(op:LazyOp) -> LazyOp: return LazyOp(op.op, tuple(get_weakop(x) if
 def get_movementroot(root:LazyBuffer) -> LazyBuffer: return get_movementroot(root.op.src[0]) if root.realized is None and (root.optype == MovementOps or (root.op.op == LoadOps.CONTIGUOUS and root.op.src[0].st.contiguous)) else root
 def get_movementroot_contiguous(x:LazyBuffer) -> LazyBuffer: return get_movementroot(x) if x.optype == MovementOps and x.st.contiguous else x
 
+def replace_with_movement_op(y:Union[LazyOp, LazyBuffer], op:MovementOps, arg:Tuple[Any, ...]) -> LazyBuffer:
+  if isinstance(y, LazyBuffer): return y.movement_op(op, arg)
+  assert y.op in BinaryOps or y.op in UnaryOps
+  return elementwise_op(y.op, *[replace_with_movement_op(z, op, arg) for z in y.src])   # type: ignore
+
+def support_weakref(x): return x
+@support_weakref  # needed for mypyc, this prevents LazyBuffer from becoming a native class
 class LazyBuffer:
   __deletable__ = ('op',)
-  lazycache : Final[weakref.WeakValueDictionary[Tuple[str, OpType, LazyOp], LazyBuffer]] = weakref.WeakValueDictionary()
+  lazycache : ClassVar[WeakValueDictionary[Tuple[str, OpType, LazyOp], LazyBuffer]] = WeakValueDictionary()
   def __new__(cls, device:str, shape:Union[ShapeTracker, Tuple[int, ...]], optype:OpType, op:LazyOp):
     # fromcpu aren't cached
     if optype == LoadOps and op.op == LoadOps.FROMCPU:
       return super().__new__(cls)
     wop = (device, optype, get_weakop(op))   # NOTE: shape should be deterministic. annoying to cache with the ShapeTracker
     # NOTE: we need "ret" to prevent the new buffer from being immediately deleted
-    if wop not in LazyBuffer.lazycache:
-      LazyBuffer.lazycache[wop] = ret = super().__new__(cls) # noqa: F841, pylint: disable=W0612
-    return LazyBuffer.lazycache[wop]
+    if wop not in LazyBuffer.lazycache: LazyBuffer.lazycache[wop] = ret = super().__new__(cls)
+    else: ret = LazyBuffer.lazycache[wop]
+    return ret
 
   def __init__(self, device:str, shape:Union[ShapeTracker, Tuple[int, ...]], optype:OpType, op:LazyOp):
     if hasattr(self, 'device'):
@@ -140,7 +142,7 @@ class LazyBuffer:
         self.realized = self.dbuffer.exec_ast(map_buffers({x:x.realize(self.device) for x in get_buffers(ast)}, ast))
       log_op(self.realized, ast)
 
-    assert self.realized.shape == self.shape
+    assert self.realized.shape == self.shape, f"shape mismatch on realize {self.realized.shape} vs {self.shape}"
     assert isinstance(self.realized, Device._buffers[self.device])
     return self.realized
 
@@ -168,7 +170,7 @@ class LazyBuffer:
     padding = tuple((max(0, -p[0]), max(0, p[1]-self.shape[i])) for i,p in enumerate(arg))
     return self.movement_op(MovementOps.PAD, padding).movement_op(MovementOps.SHRINK, tuple((p[0] + padding[i][0], p[1] + padding[i][0]) for i,p in enumerate(arg)))
 
-  def movement_op(self:LazyBuffer, op:MovementOps, arg : Tuple[Any, ...]) -> LazyBuffer:
+  def movement_op(self:LazyBuffer, op:MovementOps, arg:Tuple[Any, ...]) -> LazyBuffer:
     # very instant nop
     if op == MovementOps.RESHAPE and self.shape == arg: return self
 
@@ -200,12 +202,7 @@ class LazyBuffer:
 
     # if this MovementOp is being applied to a BinaryOp, apply the MovementOp to all the BinaryOp inputs instead
     if SHUFFLE_MOVEMENT_OPS and self.optype == BinaryOps and self.realized is None and len(self.children) == 0 and (SHUFFLE_PAD_OPS or op != MovementOps.PAD) and op not in [MovementOps.EXPAND, MovementOps.STRIDED]:
-      def replace_with_movement_op(y:Union[LazyOp, LazyBuffer]) -> LazyBuffer:
-        if isinstance(y, LazyBuffer):
-          return y.movement_op(op, arg)
-        assert y.op in BinaryOps or y.op in UnaryOps
-        return elementwise_op(y.op, *[replace_with_movement_op(z) for z in y.src])   # type: ignore
-      return replace_with_movement_op(self.op)
+      return replace_with_movement_op(self.op, op, arg)
 
     # create the buffer
     ret = LazyBuffer(self.device, ShapeTracker(self.st).movement_op(op, arg), MovementOps, LazyOp(op, (self,), arg))

tinygrad/llops/ops_gpu.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import numpy as np
-from typing import List, Tuple, Optional, Dict, Union, Set
+from typing import List, Tuple, Optional, Dict, Union, Set, Final, Callable
 from tinygrad.helpers import prod
 from tinygrad.ops import DEBUG, UnaryOps, BinaryOps, ReduceOps, MovementOps, LazyOp, Op, ExplicitExecAST, GlobalCounters
 from tinygrad.ast import ASTKernel, Token, Types
@@ -10,7 +10,7 @@ from tinygrad.shape.symbolic import ModNode   # this will go away when VALIDHACK
 from tinygrad.helpers import getenv
 
 CUDA = getenv("CUDA", 0)
-if not CUDA: from tinygrad.runtime.opencl import CLBuffer, CLImage, CLProgram, CL   # NOTE: using CL will not work for the CUDA runtime # noqa: F401
+if not CUDA: from tinygrad.runtime.opencl import CLBuffer, CLImage, CLProgram, CL, cl as pyopencl  # NOTE: using CL will not work for the CUDA runtime # noqa: F401
 else: from tinygrad.runtime.cuda import CLBuffer, CLImage, CLProgram  # type: ignore
 
 VALIDHACKS = getenv("VALIDHACKS", 0)    # TODO: remove the need for this
@@ -28,7 +28,7 @@ def split_float4(x):
   return sum([[Token(acc.tok+f".s{s}", Types.FLOAT) for s in range(4)] for acc in x], [])
 
 class CLASTKernel(ASTKernel):
-  code_for_op : Dict[Op, str] = {
+  code_for_op : Final[Dict[Op, str]] = {
     UnaryOps.NOOP: "(A)", UnaryOps.NEG: "(-(A))", UnaryOps.RELU: "max(A, (float)0.)",
     UnaryOps.GT0:  "(A > 0.)" if CUDA else "((float)1.-step((float)0.,(-A)))",
     UnaryOps.EXP: "native_exp(A)" if NATIVE_EXPLOG else "exp(A)",
@@ -38,7 +38,7 @@ class CLASTKernel(ASTKernel):
     BinaryOps.DIV: "(A/B)", BinaryOps.POW: "pow(A,B)", BinaryOps.CMPEQ: "(A==B)",
     ReduceOps.SUM: "A+=B", ReduceOps.MAX: "A=max(A,B)"
   }
-  start_for_op = {ReduceOps.SUM: "0.0", ReduceOps.MAX: "-INFINITY"}
+  start_for_op : Final[Dict[Op, str]] = {ReduceOps.SUM: "0.0", ReduceOps.MAX: "-INFINITY"}
 
   def image_idx(self, buf_index, idxy, validhacks=False):
     assert self.buftokens[buf_index].typ == Types.FLOAT4, f"image must be FLOAT4 {self.buftokens[buf_index]} {self.bufs[buf_index].st}"
@@ -223,7 +223,7 @@ class CLASTKernel(ASTKernel):
 
   # STOP WASTING TIME WITH DOING THE RESHAPES AND PERMUTES BY HAND. KERNEL SEARCH IS THE ONLY WAY IT WILL EVER BE GOOD
   # group_for_reduce will have to be better first
-  def codegen(self):
+  def codegen(self) -> Callable:
     self.process()
     self.hand_coded_optimizations()
 
@@ -257,12 +257,13 @@ class CLASTKernel(ASTKernel):
 
     # early ast
     accumulators : List[Token] = [Token("acc%d" % i, self.buftokens[0].typ) for i in range(self.buftokens[0].size())]
-    if self.reduceop:
+    if self.reduceop is not None:
       full_shape_candidates = [x.shape for x in self.sts if x.shape != self.sts[0].shape]
       full_shape : Tuple[int, ...] = self.sts[0].shape if len(full_shape_candidates) == 0 else full_shape_candidates[0]
 
       acc_offsets = self.buftokens[self.bufs.index(self.earlybufs[0])].acc_offsets()
-      self.kernel += [f"{accumulator.decltype()} {accumulator.tok} = {CLASTKernel.start_for_op[self.reduceop.op]};\n" for accumulator in accumulators]
+      assert self.reduceopop is not None
+      self.kernel += [f"{accumulator.decltype()} {accumulator.tok} = {CLASTKernel.start_for_op[self.reduceopop]};\n" for accumulator in accumulators]
       self.kernel += [f"for (int idx{i} = 0; idx{i} < {full_shape[i]}; idx{i}++) {{\n" for i in range(self.first_reduce+len(self.group_for_reduce), self.shape_len)]
       expanded_accumulators = split_float4(accumulators) if len(accumulators)*4 == len(acc_offsets) else accumulators
       self.kernel += [f"{x.tok};\n" for x in self.ast_parse(self.reduceop, [expanded_accumulators[off] for off in acc_offsets], do_reduce=True)] + ["}\n"] * (self.shape_len - (self.first_reduce + len(self.group_for_reduce)))
@@ -283,14 +284,15 @@ class CLASTKernel(ASTKernel):
         self.reshape_and_permute(None, [i for i in range(self.shape_len) if i != self.first_reduce+1] + [self.first_reduce+1])
         self.upcast()
 
+      assert self.reduceopop is not None
       self.kernel.append("if (mid_idx == 0) {\n")
       new_accumulators = [Token(f"output{i}", self.buftokens[0].typ) for i in range(len(accumulators))]
       for i,acc in enumerate(new_accumulators):
         self.kernel.append(f"{acc.decltype()} {acc.tok} = 0.0;\n")
         if self.upcast_in_mid_reduce:
-          self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)//4}; mid++) {{ {CLASTKernel.code_for_op[self.reduceop.op].replace('A', acc.tok).replace('B', f'vload4(0, &temp{i}[mid*4])')}; }}\n")
+          self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)//4}; mid++) {{ {CLASTKernel.code_for_op[self.reduceopop].replace('A', acc.tok).replace('B', f'vload4(0, &temp{i}[mid*4])')}; }}\n")
         else:
-          self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)}; mid++) {{ {CLASTKernel.code_for_op[self.reduceop.op].replace('A', acc.tok).replace('B', f'temp{i}[mid]')}; }}\n")
+          self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)}; mid++) {{ {CLASTKernel.code_for_op[self.reduceopop].replace('A', acc.tok).replace('B', f'temp{i}[mid]')}; }}\n")
       accumulators = new_accumulators
     
     # late ast
@@ -326,17 +328,19 @@ class CLASTKernel(ASTKernel):
 class GPUBuffer(ExplicitExecAST):
   def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf:Optional[GPUBuffer]=None, backing:Optional[np.ndarray]=None, force_create=False):
     super().__init__(shape, hostbuf)
-    self._buf : Optional[CLBuffer] = hostbuf._buf if hostbuf is not None else None
+    self._buf : Optional[Union[CLImage, CLBuffer]] = hostbuf._buf if hostbuf is not None else None
     self._base_shape : Tuple[int, ...] = hostbuf._base_shape if hostbuf is not None else self.shape
     self._backing : Optional[np.ndarray] = hostbuf._backing if hostbuf is not None else backing
     # early copy in for large buffers
     if (self._backing is not None and self._backing.shape != (1,)) or force_create:
       self.cl
   
+  # TODO: refactor this to return self._buf and not import pyopencl
   @property
-  def cl(self):
+  def cl(self) -> pyopencl.Buffer:
     if self._buf is None:
       self._buf = CLImage(self._base_shape) if (len(self._base_shape) == 3 and self._base_shape[2] == 4 and IMAGE >= 2) else CLBuffer(4*prod(self._base_shape))
+    assert self._buf is not None
     if self._backing is not None:
       self._buf.copyin(self._backing)
       self._backing = None

tinygrad/ops.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import numpy as np
-from enum import Enum
-from typing import Union, Type, NamedTuple, Tuple, Any, List
+from enum import Enum, auto
+from typing import Union, Type, NamedTuple, Tuple, Any, List, ClassVar
 import functools, operator
 from tinygrad.helpers import prod
 from tinygrad.shape import ShapeTracker
@@ -10,12 +10,13 @@ from tinygrad.helpers import getenv
 DEBUG = getenv("DEBUG", 0)
 
 # these are the llops your accelerator must implement, along with toCpu
-UnaryOps = Enum("UnaryOps", ["NOOP", "NEG", "RELU", "EXP", "LOG", "GT0", "RECIPROCAL"])
-BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "CMPEQ"])
-ReduceOps = Enum("ReduceOps", ["SUM", "MAX"])
-MovementOps = Enum("MovementOps", ["RESHAPE", "PERMUTE", "EXPAND", "FLIP", "STRIDED", "PAD", "SHRINK"])
-ProcessingOps = Enum("ProcessingOps", ["CONV"])
-LoadOps = Enum("LoadOps", ["FROMCPU", "CONTIGUOUS"])
+# the Enum class doesn't work with mypy, this is static. sorry it's ugly
+class UnaryOps(Enum): NOOP = auto(); NEG = auto(); RELU = auto(); EXP = auto(); LOG = auto(); GT0 = auto(); RECIPROCAL = auto() # noqa: E702
+class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto() # noqa: E702
+class ReduceOps(Enum): SUM = auto(); MAX = auto() # noqa: E702
+class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); FLIP = auto(); STRIDED = auto(); PAD = auto(); SHRINK = auto() # noqa: E702
+class ProcessingOps(Enum): CONV = auto() # noqa: E702
+class LoadOps(Enum): FROMCPU = auto(); CONTIGUOUS = auto() # noqa: E702
 
 Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps]
 OpType = Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOps], Type[ProcessingOps], Type[LoadOps]]
@@ -30,6 +31,9 @@ class LazyOp(NamedTuple):
 # Any == Union[LazyBuffer, DeviceBuffer]
 def get_buffers(op:LazyOp) -> List[Any]: return functools.reduce(operator.add, [get_buffers(x) if isinstance(x, LazyOp) else [x] for x in op.src], [])
 def get_lazyops(op:LazyOp) -> List[LazyOp]: return functools.reduce(operator.add, [get_lazyops(x) for x in op.src if isinstance(x, LazyOp)], [op])
+def map_buffers(real_srcs, x:LazyOp) -> LazyOp:
+  if x in real_srcs: return map_buffers(real_srcs, real_srcs[x]) if isinstance(real_srcs[x], LazyOp) else real_srcs[x]
+  return LazyOp(x.op, tuple((map_buffers(real_srcs, y) if isinstance(y, LazyOp) else real_srcs[y]) for y in x.src), x.arg)
 
 # a placeholder class to extend by the exec classes
 class DeviceBuffer:
@@ -63,7 +67,10 @@ class GenericExecAST(DeviceBuffer):  # pylint: disable=abstract-method
     return ret
 
 class GlobalCounters:
-  global_ops, global_mem, time_sum = 0, 0, 0
+  global_ops : ClassVar[int] = 0
+  global_mem : ClassVar[int] = 0
+  time_sum : ClassVar[int] = 0
+  # TODO: reset method
 
 class GenericShape(GenericExecAST):  # pylint: disable=abstract-method
   def __init__(self, shape, flops=0): self.shape, self.flops = shape, flops

tinygrad/runtime/opencl.py

@@ -58,6 +58,9 @@ class CLImage:
   def __del__(self):
     CL.mem_used -= self.cl.row_pitch * self.cl.height
 
+  def copyin(self, b:np.ndarray): raise NotImplementedError("no copyin for CLImage")
+  def copyout(self, a:np.ndarray): raise NotImplementedError("no copyout for CLImage")
+
 @functools.lru_cache(maxsize=None)
 class CLProgram:
   kernel_cnt : Dict[str, int] = defaultdict(int)

tinygrad/tensor.py

@@ -3,13 +3,15 @@ from __future__ import annotations
 import inspect, functools, importlib, itertools
 import numpy as np
 from tinygrad.helpers import prod, argfix, make_pair
-from typing import List, Tuple, Callable, Optional
+from typing import List, Tuple, Callable, Optional, ClassVar, Type
 from tinygrad.lazy import Device, LazyBuffer
 
 # **** start with two base classes, Tensor and Function ****
 
 class Tensor:
-  training, no_grad = False, False
+  __deletable__ = ('_ctx',)
+  training : ClassVar[bool] = False
+  no_grad : ClassVar[bool] = False
 
   def __init__(self, data, device=Device.DEFAULT, requires_grad:Optional[bool]=None):
     if isinstance(data, list):
@@ -40,14 +42,15 @@ class Tensor:
     return f"<Tensor {self.lazydata if self.lazydata.realized is None else self.lazydata.realized!r} with grad {(self.grad.lazydata if self.grad else None)!r}>"
 
   @property
-  def shape(self): return self.lazydata.shape
+  def shape(self) -> Tuple[int, ...]: return self.lazydata.shape
+
 
   # dtype handling was very broken. it's always float32 now
   @property
-  def dtype(self): return np.float32
+  def dtype(self) -> type: return np.float32
 
   @property
-  def device(self): return self.lazydata.device
+  def device(self) -> str: return self.lazydata.device
 
   # ***** data handlers ****
 
@@ -63,11 +66,11 @@ class Tensor:
     return x
 
   def detach(self): return Tensor(self.lazydata, device=self.device, requires_grad=False)
-  def numpy(self): return np.array(self.lazydata.toCPU())
+  def numpy(self) -> np.ndarray: return np.array(self.lazydata.toCPU())
 
   # TODO: this keeps the legacy behavior working, remove it after refactor
   @property
-  def data(self): return self.numpy()
+  def data(self) -> np.ndarray: return self.numpy()
 
   # TODO: if things are realized this won't work
   def to_(self, device:str):
@@ -168,11 +171,11 @@ class Tensor:
     val = [val] if not isinstance(val, (list, tuple)) else val
     assert sum(s is not None for s in val) <= len(self.shape)
     assert all(s.step is None or s.step == 1 for s in val if isinstance(s, slice))
-    for i,(sz,s) in enumerate(zip(self.shape, (v for v in val if v is not None))):  # Slicing only depends on ints + slices
+    for i,(sz,s) in enumerate(zip(self.shape, [v for v in val if v is not None])):  # Slicing only depends on ints + slices
       if isinstance(s, int) and not (-sz <= s < sz):
         raise IndexError(f"index {s} is out of bounds for dimension {i} with size {sz}")
       new_slice.append((s%sz, s%sz+1) if isinstance(s, int) else (slcfix(s.start, sz, 0), slcfix(s.stop, sz, sz)))
-    for s,sz in zip(val, (self.shape[i-1] for i in itertools.accumulate(s is not None for s in val))):  # Shape depends on slices + positions of Nones
+    for s,sz in zip(val, [self.shape[i-1] for i in itertools.accumulate([s is not None for s in val])]):  # Shape depends on slices + positions of Nones
       if not isinstance(s, int):
         new_shape.append(1 if s is None else slcfix(s.stop, sz, sz) - slcfix(s.start, sz, 0))
     new_shape += [self.shape[i] for i in range(len(new_slice), len(self.shape))]
@@ -184,7 +187,7 @@ class Tensor:
     for y in args:
       assert len(y.shape) == len(self.shape) and all(y.shape[i] == s for i,s in enumerate(self.shape) if i != dim)
     catargs = [self] + list(args)
-    shape_cumsum = [0, *itertools.accumulate(y.shape[dim] for y in catargs)]
+    shape_cumsum = [0, *itertools.accumulate([y.shape[dim] for y in catargs])]
     slc = [[(0, s) for s in self.shape] for _ in catargs]
     for s,k in zip(slc, shape_cumsum):
       s[dim] = (-k, shape_cumsum[-1]-k)
@@ -216,7 +219,7 @@ class Tensor:
     return cx.conv2d(cw, groups=groups).reshape(shape=out_shape_t).transpose(order=order)
 
   # TODO: what's the difference between dot and matmul?
-  dot = matmul
+  def dot(self:Tensor, w:Tensor): return self.matmul(w)
 
   # (padding_left, padding_right, padding_top, padding_bottom)
   def pad2d(self, padding:Tuple[int, ...]): return self.slice(arg = [(0,self.shape[0]), (0,self.shape[1]), (-padding[2],self.shape[2]+padding[3]), (-padding[0],self.shape[3]+padding[1])])  # type: ignore
@@ -287,7 +290,7 @@ class Tensor:
   def sigmoid(self): return (1.0 + (-self).exp()).reciprocal()
   def elu(self, alpha=1.0): return self.relu() - alpha*(1-self.exp()).relu()
   def swish(self): return self * self.sigmoid()
-  silu = swish   # The SiLU function is also known as the swish function.
+  def silu(self): return self.swish()   # The SiLU function is also known as the swish function.
   def relu6(self): return self.relu() - (self-6).relu()
   def hardswish(self): return self * (self+3).relu6() * (1/6)
   def tanh(self): return 2.0 * ((2.0 * self).sigmoid()) - 1.0
@@ -307,12 +310,15 @@ class Tensor:
     shape_ret = tuple(max(sx, sy) for sx,sy in zip(x.shape, y.shape))
     return fxn(x.expand(shape_ret), y.expand(shape_ret))
 
+  @staticmethod
+  def _truediv(x, y): return x * (y.reciprocal() if isinstance(y, Tensor) else (1/y))
+
   # TODO: are these the only ones that can take number arguments?
   def add(self, x): return Tensor.broadcasted(Tensor._add, self, x)
   def sub(self, x): return Tensor.broadcasted(Tensor._sub, self, x)
   def mul(self, x): return Tensor.broadcasted(Tensor._mul, self, x)
   def pow(self, x): return Tensor.broadcasted(Tensor._pow, self, x)
-  def div(self, y): return self * (y.reciprocal() if isinstance(y, Tensor) else (1/y))
+  def div(self, x): return Tensor._truediv(self, x)
 
   # ***** functional nn ops *****
 
@@ -349,30 +355,32 @@ class Function:
   # NOTE: it doesn't hurt to save this since the ctx will be freed fast without grad
   def save_for_backward(self, *x): self.saved_tensors.extend(x)
 
-  @classmethod
-  def apply(cls, *x:Tensor, **kwargs):
-    ctx = cls(x[0].device, *x)
+  @staticmethod
+  def apply(fxn:Type[Function], *x:Tensor, **kwargs):
+    ctx = fxn(x[0].device, *x)
     ret = Tensor(ctx.forward(*[t.lazydata for t in x], **kwargs), device=ctx.device, requires_grad=ctx.requires_grad)
     if ctx.requires_grad and not Tensor.no_grad:
       ret._ctx = ctx    # used by autograd engine
     return ret
 
 # register functions to move between devices
-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]))
-  setattr(Tensor, f"{device.lower()}_", functools.partialmethod(Tensor.to_, Device.__dict__[device]))
+for device in [device for device in Device._buffers.keys() if device[0] != "_"]:
+  setattr(Tensor, f"{device.lower()}", functools.partialmethod(Tensor.to, device))
+  setattr(Tensor, f"{device.lower()}_", functools.partialmethod(Tensor.to_, device))
 
 # register all the mlops "math" operations
-def register(name:str, fxn:Function):
-  setattr(Tensor, "_"+name if hasattr(Tensor, name) else name, lambda *args, **kwargs: fxn.apply(*args, **kwargs))  # doesn't work without lambda, pylint: disable=W0108
+def register(name:str, fxn:Type[Function]):
+  setattr(Tensor, "_"+name if hasattr(Tensor, name) else name, lambda *args, **kwargs: Function.apply(fxn, *args, **kwargs))  # doesn't work without lambda, pylint: disable=W0108
 for name, cls in inspect.getmembers(importlib.import_module('tinygrad.mlops'), inspect.isclass):
   if name[0] != "_" and name != "Function" and not name.endswith("Ops"):
     register(name.lower(), cls)
 
 # register the operators
-def register_op(name, fxn):
+def register_op(name, fop):
+  if name in ['add', 'sub', 'mul', 'pow']: fxn = lambda x,y: Tensor.broadcasted(fop, x, y)
+  else: fxn = lambda x,y: fop(x,y)  # pylint: disable=W0108
   setattr(Tensor, f"__{name}__", fxn)
-  setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(fxn(self,x)))
-  setattr(Tensor, f"__r{name}__", lambda self,x: fxn(x,self))
-for name in ['add', 'sub', 'mul', 'pow', 'matmul', 'truediv']:
-  register_op(name, getattr(Tensor, name if name != 'truediv' else 'div'))
+  setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(fxn(self, x)))
+  setattr(Tensor, f"__r{name}__", lambda self,x: fxn(x, self))
+for name in ['add', 'sub', 'mul', 'pow', "truediv", "matmul"]:
+  register_op(name, getattr(Tensor, ("_"+name) if name != "matmul" else name))