sgd now operates on Tensors
MODIFY
tinygrad/optim.py
tinygrad/optim.py
# sorted in order of increasing complexity
import numpy as np
+from tinygrad.tensor import Tensor
class Optimizer:
def __init__(self, params):
@@ -9,11 +10,11 @@ class Optimizer:
class SGD(Optimizer):
def __init__(self, params, lr=0.001):
super(SGD, self).__init__(params)
- self.lr = lr
+ self.lr = Tensor([lr])
def step(self):
for t in self.params:
- t.data -= self.lr * t.grad.data
+ t -= t.grad * self.lr
class RMSprop(Optimizer):
def __init__(self, params, lr=0.001, decay=0.9, eps=1e-8):

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -47,6 +47,9 @@ class Tensor:
def __repr__(self):
return "Tensor %r with grad %r" % (self.data, self.grad.data if self.grad else None)
+ def assign(self, x):
+ self.data = x.data
+
@property
def shape(self):
return self.data.shape
@@ -125,6 +128,32 @@ class Tensor:
ops = {}
opsgpu = {}
+ # ***** basic ops *****
+
+ def __add__(self, x):
+ return self.add(x)
+
+ def __sub__(self, x):
+ return self.sub(x)
+
+ def __mul__(self, x):
+ return self.mul(x)
+
+ def __div__(self, x):
+ return self.div(x)
+
+ def __iadd__(self, x):
+ self.assign(self + x)
+
+ def __isub__(self, x):
+ self.assign(self - x)
+
+ def __imul__(self, x):
+ self.assign(self * x)
+
+ def __idiv__(self, x):
+ self.assign(self / x)
+
# ***** non first class ops *****
def mean(self):

sgd now operates on Tensors