wooo 1k lines
MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -203,11 +203,8 @@ class LazyBuffer:
return self.realized
@staticmethod
- def fromCPU(x, device):
- return LazyBuffer(device, x.shape, LoadOps, LazyOp(LoadOps.FROMCPU, tuple(), x.copy()))
-
- def toCPU(x):
- return x.realize().toCPU()
+ def fromCPU(x, device): return LazyBuffer(device, x.shape, LoadOps, LazyOp(LoadOps.FROMCPU, tuple(), x.copy()))
+ def toCPU(x): return x.realize().toCPU()
def unary_op(x:LazyBuffer, op:UnaryOps) -> LazyBuffer: return elementwise_op(op, x)
def binary_op(x:LazyBuffer, op:BinaryOps, y:LazyBuffer) -> LazyBuffer: return elementwise_op(op, x, y)

MODIFY
tinygrad/shapetracker.py
tinygrad/shapetracker.py
@@ -75,8 +75,7 @@ class ShapeTracker:
self.views : List[ViewTypes] = shape.views[:] if isinstance(shape, ShapeTracker) else [view_from_shape(shape)]
@property
- def contiguous(self):
- return len(self.views) == 1 and self.views[-1].contiguous
+ def contiguous(self): return len(self.views) == 1 and self.views[-1].contiguous
@property
def shape(self): return self.views[-1].shape
@@ -119,8 +118,7 @@ class ShapeTracker:
# TODO: this is a special case of slice with strides, remove it
# though it's nice that it can't change size
- def flip(self, *axis):
- self.stride(*[-1 if i in axis else 1 for i in range(len((self.shape)))])
+ def flip(self, *axis): self.stride(*[-1 if i in axis else 1 for i in range(len((self.shape)))])
# *** under this line are not invertible ***

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -331,12 +331,12 @@ class Tensor:
# An instantiation of the Function is the Context
class Function:
- def __init__(self, device, *tensors):
+ def __init__(self, device, *tensors:Tensor):
self.device = device
self.parents = tensors
self.needs_input_grad = [t.requires_grad for t in tensors]
self.requires_grad = any(self.needs_input_grad) and not Tensor.no_grad
- self.saved_tensors = []
+ self.saved_tensors : List[Tensor] = []
def forward(self, *args, **kwargs): raise NotImplementedError(f"forward not implemented for {type(self)}")
def backward(self, *args, **kwargs): raise NotImplementedError(f"backward not implemented for {type(self)}")
@@ -348,8 +348,7 @@ class Function:
@classmethod
def apply(cls, *x:Tensor, **kwargs):
ctx = cls(x[0].device, *x)
- ret = Tensor(ctx.forward(*[t.lazydata for t in x], **kwargs),
- device=ctx.device, requires_grad=ctx.requires_grad)
+ ret = Tensor(ctx.forward(*[t.lazydata for t in x], **kwargs), device=ctx.device, requires_grad=ctx.requires_grad)
if ctx.requires_grad: ret._ctx = ctx # used by autograd engine
return ret

wooo 1k lines