cleanups to Tensor class
MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -26,7 +26,7 @@ class ReLU(_UnaryOp):
class Log(_UnaryOp):
fop = UnaryOps.LOG
- bop = BinaryOps.DIV
+ bop = BinaryOps.DIV # TODO: flip order of DIV
class Exp(_UnaryOp):
def forward(ctx, input):
@@ -36,6 +36,8 @@ class Exp(_UnaryOp):
bop = BinaryOps.MUL
+# TODO: add Neg?
+
# ************* reduce ops *************
class Sum(Function):
@@ -96,6 +98,8 @@ class Mul(Function):
grad_y = ctx.binary_op(BinaryOps.MUL, x, grad_output) if ctx.needs_input_grad[1] else None
return grad_x, grad_y
+# TODO: add Div?
+
class Pow(Function):
def forward(ctx, x, y):
ret = ctx.binary_op(BinaryOps.POW, x, y)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -41,6 +41,7 @@ class Tensor:
return f"<Tensor {self.data!r} with grad {(self.grad.data if self.grad else None)!r}>"
def realize(self):
+ # TODO: once lazy is upstreamed, we can remove this check
if getattr(self.data, 'realize', None) is not None:
self.data.realize()
@@ -65,6 +66,8 @@ class Tensor:
# ***** creation helper functions *****
+ # TODO: remove use of numpy here
+
@classmethod
def zeros(cls, *shape, **kwargs):
return cls(np.zeros(shape, dtype=np.float32), **kwargs)
@@ -155,7 +158,7 @@ class Tensor:
def numpy(self):
return np.array(self.cpu().data)
- # ***** non first class ops *****
+ # ***** non first class ops (hlops) *****
def __getitem__(self, val):
arg = []
@@ -213,6 +216,7 @@ class Tensor:
cw = w.transpose(order=worder).reshape(shape=(groups*cout, cin, 1, 1))
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 transpose(self, order=(1,0)):
@@ -243,43 +247,6 @@ class Tensor:
out = self.sum(axis=axis, keepdim=keepdim)
return out * (prod(out.shape)/prod(self.shape))
- def sqrt(self):
- return self.pow(0.5)
-
- def div(self, y):
- return self * (y ** -1.0)
- __truediv__ = div
-
- def sigmoid(self):
- #e = self.exp(); return e.div(1 + e)
- return (1.0 + (0.0-self).exp()) ** -1.0
-
- def elu(self, alpha=1.0):
- return self.relu() - (-alpha*(self.exp() - 1)).relu()
-
- def swish(self):
- return self * self.sigmoid()
-
- def relu6(self):
- return self.relu() - (self-6).relu()
-
- def clip(self, min, max):
- return ((self-min).relu()+min) - (self-max).relu()
-
- def hardswish(self):
- return self * (self+3).relu6() * (1/6)
-
- def tanh(self):
- return 2.0 * ((2.0 * self).sigmoid()) - 1.0
-
- def gelu(x):
- # https://github.com/huggingface/transformers/blob/master/src/transformers/activations.py
- #import torch; return Tensor(torch.nn.functional.gelu(torch.tensor(x.data)).numpy())
- return 0.5 * x * (1 + (x * 0.7978845608 * (1 + 0.044715 * x * x)).tanh())
-
- def leakyrelu(self, neg_slope=0.01):
- return self.relu() - (-neg_slope*self).relu()
-
def _softmax(self):
m = self - self.max(axis=len(self.shape)-1, keepdim=True)
e = m.exp()
@@ -298,21 +265,7 @@ class Tensor:
_mask = np.asarray(np.random.binomial(1, 1.0-p, size=self.shape), dtype=self.dtype)
return self * Tensor(_mask, requires_grad=False, device=self.device) * (1/(1.0 - p))
- def softplus(self, limit=20, beta=1):
- # safe softplus - 1/beta*log(1 + exp(beta*x)) (PyTorch)
- eb = (self*beta).exp()
- ret = (1 + eb).log()
- return (1/beta)*ret
-
- def mish(self):
- return self * (self.softplus().tanh()) # x*tanh(softplus(x))
-
- def abs(self):
- return self.relu() + (-1.0*self).relu()
-
- def sign(self):
- return self / (self.abs() + 1e-10)
-
+ # TODO: support arbitrary strides
def _pool2d(self, py, px):
xup = self[:, :, :self.shape[2]-self.shape[2]%py, :self.shape[3]-self.shape[3]%px] if (self.shape[2]%py != 0) or (self.shape[3]%px != 0) else self
return xup.reshape(shape=(xup.shape[0], xup.shape[1], xup.shape[2]//py, py, xup.shape[3]//px, px))
@@ -327,6 +280,27 @@ class Tensor:
ret = self._conv2d(weight, **kwargs)
return ret if bias is None else ret.add(bias.reshape(shape=[1, -1, 1, 1]))
+ # ***** math functions (unary) *****
+
+ def sqrt(self): return self.pow(0.5)
+ def clip(self, min, max): return ((self-min).relu()+min) - (self-max).relu()
+ def abs(self): return self.relu() + (0.0-self).relu()
+ def sign(self): return self / (self.abs() + 1e-10)
+
+ # ***** activation functions (unary) *****
+
+ # TODO: make "-self" work
+ def sigmoid(self): return (1.0 + (0.0-self).exp()) ** -1.0
+ def elu(self, alpha=1.0): return self.relu() - (-alpha*(self.exp() - 1)).relu()
+ def swish(self): return self * self.sigmoid()
+ 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
+ def gelu(x): return 0.5 * x * (1 + (x * 0.7978845608 * (1 + 0.044715 * x * x)).tanh())
+ def leakyrelu(self, neg_slope=0.01): return self.relu() - (-neg_slope*self).relu()
+ def mish(self): return self * self.softplus().tanh()
+ def softplus(self, limit=20, beta=1): return (1/beta) * (1 + (self*beta).exp()).log()
+
# ***** broadcasted binary ops *****
@staticmethod
@@ -350,9 +324,14 @@ class Tensor:
def mul(self, x): return Tensor.broadcasted(Tensor._mul, self, x)
def pow(self, x): return Tensor.broadcasted(Tensor._pow, self, x)
+ # TODO: should be broadcasted binary op
+ def div(self, y):
+ return self * (y ** -1.0)
+ __truediv__ = div
+
# ***** functional nn ops *****
- # TODO: fix the kwargs problem
+ # TODO: fix the kwargs problem, then remove these
def reshape(self, shape): return self._reshape(shape=shape)
def expand(self, shape): return self._expand(shape=shape)
@@ -390,8 +369,7 @@ class Function(Ops):
ctx = cls(x[0].device, *x)
ret = Tensor(ctx.forward(*[t.data for t in x], **kwargs),
device=ctx.device, requires_grad=ctx.requires_grad)
- if ctx.requires_grad:
- ret._ctx = ctx # used by autograd engine
+ if ctx.requires_grad: ret._ctx = ctx # used by autograd engine
return ret
# register functions to move between devices

cleanups to Tensor class