minor optimizations & cleaning (#257)

* use isinstance, some optimizations & whitespace removal

* revert whitespace changes

* revert more whitespace

* some more cleanup

* revert fstring (not a fan of the {{}})

* fix typo

* fix typo

examples/serious_mnist.py

@@ -135,4 +135,4 @@ if __name__ == "__main__":
       X_aug = X_train if epoch == 1 else augment_img(X_train)
       train(model, X_aug, Y_train, optimizer, steps=steps, lossfn=lossfn, BS=BS)
       accuracy = evaluate(model, X_test, Y_test, BS=BS)
-      model.save('examples/checkpoint'+str("%.0f" % (accuracy*1.0e6)))
+      model.save(f'examples/checkpoint{accuracy * 1e6:.0f}')

examples/yolo_nn.py

@@ -5,8 +5,7 @@ from tinygrad.tensor import Tensor
 
 class MaxPool2d:
   def __init__(self, kernel_size, stride):
-    if type(kernel_size) == int:
-      self.kernel_size = (kernel_size, kernel_size)
+    if isinstance(kernel_size, int): self.kernel_size = (kernel_size, kernel_size)
     else: self.kernel_size = kernel_size
     self.stride = stride if (stride is not None) else kernel_size
   
@@ -62,9 +61,7 @@ class Conv2d:
   def __init__(self, in_channels, out_channels, kernel_size, stride = 1, padding = 0, groups = 1, bias = True):
     self.in_channels, self.out_channels, self.stride, self.padding, self.groups, self.bias = in_channels, out_channels, stride, padding, groups, bias # Wow this is terrible
 
-    if type(kernel_size) == int:
-      self.kernel_size = (kernel_size, kernel_size)
-    else: self.kernel_size = kernel_size
+    self.kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
 
     assert out_channels % groups == 0 and in_channels % groups == 0
 

examples/yolov3.py

@@ -64,7 +64,7 @@ def letterbox_image(img, inp_dim=608):
   return canvas
 
 def add_boxes(img, prediction):
-  if type(prediction) is int: # no predictions
+  if isinstance(prediction, int): # no predictions
     return img
   coco_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names')
   coco_labels = coco_labels.decode('utf-8').split('\n')

extra/augment.py

@@ -37,4 +37,4 @@ if __name__ == "__main__":
   for i in range(10):
     im = Image.fromarray(X_train[7353+i])
     im_aug = [Image.fromarray(x) for x in augment_img(np.array([X_train[7353+i]]*100))]
-    im.save("aug"+str(i)+".gif", save_all=True, append_images=im_aug, duration=100, loop=0)
+    im.save(f"aug{i}.gif", save_all=True, append_images=im_aug, duration=100, loop=0)

extra/kinne.py

@@ -51,10 +51,10 @@ class KinneDir:
     """
     parameter loads or saves a parameter, given as a tensor.
     """
-    path = self.base + str(self.next_part_index) + ".bin"
+    path = f"{self.base}{self.next_part_index}.bin"
     if self.save:
       t.data.astype("<f4", "C").tofile(path)
-      self.metadata.write(str(self.next_part_index) + ": " + str(t.shape) + "\n")
+      self.metadata.write(f"{self.next_part_index}: {t.shape}\n")
     else:
       t.assign(Tensor(numpy.fromfile(path, "<f4")).reshape(shape=t.shape))
     self.next_part_index += 1
@@ -73,4 +73,3 @@ class KinneDir:
   def close(self):
     if self.save:
       self.metadata.close()
-

extra/utils.py

@@ -24,7 +24,7 @@ def get_parameters(obj):
     for x in obj:
       parameters.extend(get_parameters(x))
   elif hasattr(obj, '__dict__'):
-    for k,v in obj.__dict__.items():
+    for v in obj.__dict__.values():
       parameters.extend(get_parameters(v))
   return parameters
 

tinygrad/ops_cpu.py

@@ -1,5 +1,5 @@
 import numpy as np
-from .tensor import Function, register
+from .tensor import Function
 
 # ************* unary ops *************
 
@@ -40,13 +40,13 @@ class Sum(Function):
 
   def backward(ctx, grad_output):
     input, axis = ctx.saved_tensors
-    axis = [axis] if type(axis) is int else axis
+    if isinstance(axis, int): axis = [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)
 
 class Max(Function):
   def forward(ctx, inp, axis=None):
-    axis = [axis] if type(axis) == int else axis
+    if isinstance(axis, int): axis = [axis]
     ret = np.amax(inp, axis=None if axis is None else tuple(axis), keepdims=True)
     ctx.save_for_backward(inp, axis, ret)
     if axis is not None:
@@ -154,8 +154,7 @@ class Matmul(Function):
 
 class Conv2D(Function):
   def forward(ctx, x, w, stride=1, groups=1):
-    if type(ctx.stride) == int:
-      ctx.stride = (ctx.stride, ctx.stride)
+    if isinstance(ctx.stride, int): ctx.stride = (ctx.stride, ctx.stride)
     cout,cin,H,W = w.shape
     ys,xs = ctx.stride
     bs,cin_ = x.shape[0], x.shape[1]

tinygrad/ops_gpu.py

@@ -1,7 +1,7 @@
 import functools
 import pyopencl as cl
 import numpy as np
-from .tensor import Function, register, GPUBuffer, Tensor, Device
+from .tensor import Function, GPUBuffer
 
 def buffer_new(ctx, shape, zero=False):
   return GPUBuffer(shape, hostbuf=None if not zero else np.zeros(shape, dtype=np.float32))
@@ -9,7 +9,7 @@ def buffer_new(ctx, shape, zero=False):
 def buffer_np(ctx, x):
   return cl.Buffer(ctx.cl_ctx, cl.mem_flags.READ_WRITE | cl.mem_flags.COPY_HOST_PTR, hostbuf=x)
 
-@functools.lru_cache()
+@functools.lru_cache
 def clbuild(cl_ctx, name, prg):
   return cl.Program(cl_ctx, prg).build().__getattr__(name)
 
@@ -106,7 +106,7 @@ def reduce_op(ctx, code, code2, inp, axis=None, start="0.0"):
 
 class Sum(Function):
   def forward(ctx, input, axis=None):
-    axis = [axis] if type(axis) == int else axis
+    if isinstance(axis, int): axis = [axis]
     ctx.save_for_backward(input, axis)
     ret = reduce_op(ctx, "out += a", "out", input, axis=axis)
     if axis is not None:
@@ -121,7 +121,7 @@ class Sum(Function):
 
 class Max(Function):
   def forward(ctx, input, axis=None):
-    axis = [axis] if type(axis) == int else axis
+    if isinstance(axis, int): axis = [axis]
     ret = reduce_op(ctx, "out = max(a,out)", "out", input, axis=axis, start="-INFINITY")
     ctx.save_for_backward(input, axis, ret)
     if axis is not None:
@@ -138,11 +138,11 @@ class Max(Function):
 
 # ************* binary ops *************
 
-@functools.lru_cache()
+@functools.lru_cache
 def get_binop_prg(cl_ctx, code, complist):
   ndims = len(complist)
-  args = "".join([", int d%d" % i for i in range(ndims)]) + "".join([", int p%d" % i for i in range(ndims-1)])
-  compute_idx_rets = ["\n    int idx_ret"+str(i)+" = (gid0 / "+("p%d"%i if i < ndims-1 else "1")+") % d"+str(i)+";" for i in range(ndims)]
+  args = "".join([f", int d{i}" for i in range(ndims)] + [f", int p{i}" for i in range(ndims-1)])
+  compute_idx_rets = "".join([f"\n    int idx_ret{i} = (gid0 / {f'p{i}' if i < ndims-1 else '1'}) % d{i};" for i in range(ndims)])
 
   idx_exprs = ["0", "0"] # [idx_x, idx_y]
   for i in range(ndims):
@@ -151,7 +151,7 @@ def get_binop_prg(cl_ctx, code, complist):
         idx_exprs[j] = "idx_ret%d + d%d*(%s)" % (i, i, idx_exprs[j])
 
   return cl.Program(cl_ctx, """__kernel void binop(__global const float *x_g, __global const float *y_g, __global float *res_g"""+args+""") {
-    int gid0 = get_global_id(0);"""+"".join(compute_idx_rets)+"""
+    int gid0 = get_global_id(0);"""+compute_idx_rets+"""
     float a = x_g["""+idx_exprs[0]+"""];
     float b = y_g["""+idx_exprs[1]+"""];
     res_g[gid0] = """+code+""";\n}""").build()
@@ -365,8 +365,7 @@ class Matmul(Function):
 
 class Conv2D(Function):
   def forward(ctx, x, w, stride=1, groups=1):
-    if type(ctx.stride) == int:
-      ctx.stride = (ctx.stride, ctx.stride)
+    if isinstance(ctx.stride, int): ctx.stride = (ctx.stride, ctx.stride)
     cout,cin,H,W = w.shape
     ys,xs = ctx.stride
     bs,cin_,iy,ix = x.shape

tinygrad/tensor.py

@@ -1,5 +1,4 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
-import sys
 import inspect
 import functools
 import os
@@ -208,13 +207,15 @@ class Tensor:
   
   def __getitem__(self, val):
     arg = []
-    for i,s in enumerate(val if type(val) in [list, tuple] else ([] if val is None else [val])):
-      if type(s) is int: arg.append((s, s+1))
-      else:
-        arg.append((s.start if s.start is not None else 0,
-          (s.stop if s.stop >=0 else self.shape[i]+s.stop) if s.stop is not None else self.shape[i]))
-      if type(s) is not int: assert s.step is None or s.step == 1
-    return self.slice(arg = arg+[(0,self.shape[i]) for i in range(len(arg), len(self.shape))])
+    if val is not None:
+      for i, s in enumerate(val if isinstance(val, (list, tuple)) else [val]):
+        if isinstance(s, int):
+          arg.append((s, s + 1))
+        else:
+          arg.append((s.start if s.start is not None else 0,
+            (s.stop if s.stop >=0 else self.shape[i]+s.stop) if s.stop is not None else self.shape[i]))
+          assert s.step is None or s.step == 1
+    return self.slice(arg = arg + [(0,self.shape[i]) for i in range(len(arg), len(self.shape))])
 
   def pad2d(self, padding):
     return self[:, :, -padding[2]:self.shape[2]+padding[3], -padding[0]:self.shape[3]+padding[1]]
@@ -303,7 +304,7 @@ class Function:
   def __new__(cls, *args, **kwargs):
     cls.forward = staticmethod(cls.forward)
     cls.backward = staticmethod(cls.backward)
-    return super().__new__(cls) #
+    return super().__new__(cls)
 
   def __init__(self, *tensors):
     self.parents = tensors