more similar between torch and numpy
MODIFY
tinygrad/ops_cpu.py
tinygrad/ops_cpu.py
@@ -6,7 +6,7 @@ from .tensor import Function
class ReLU(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
- return np.maximum(input, 0)
+ return input.relu()
def backward(ctx, grad_output):
input, = ctx.saved_tensors
@@ -15,7 +15,7 @@ class ReLU(Function):
class Log(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
- return np.log(input)
+ return input.log()
def backward(ctx, grad_output):
input, = ctx.saved_tensors
@@ -23,7 +23,7 @@ class Log(Function):
class Exp(Function):
def forward(ctx, input):
- ret = np.exp(input)
+ ret = input.exp()
ctx.save_for_backward(ret)
return ret
@@ -36,7 +36,7 @@ class Exp(Function):
class Sum(Function):
def forward(ctx, input, axis=None):
ctx.save_for_backward(input, axis)
- return np.array([input.sum()]) if axis is None else input.sum(axis=axis)
+ return input.sum(axis) if axis != None else input.sum().reshape((1,))
def backward(ctx, grad_output):
input, axis = ctx.saved_tensors
@@ -105,7 +105,7 @@ class Pow(Function):
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
return unbroadcast(y * (x**(y-1.0)) * grad_output, x.shape), \
- unbroadcast((x**y) * np.log(x) * grad_output, y.shape)
+ unbroadcast((x**y) * x.log() * grad_output, y.shape)
# ************* movement ops *************

MODIFY
tinygrad/ops_torch.py
tinygrad/ops_torch.py
@@ -2,49 +2,9 @@ import torch
import numpy as np
from .tensor import Function
-# ************* unary ops *************
+# ************* unary+binary ops *************
-class ReLU(Function):
- def forward(ctx, input):
- ctx.save_for_backward(input)
- return input.relu()
-
- def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- return grad_output * (input >= 0)
-
-class Log(Function):
- def forward(ctx, input):
- ctx.save_for_backward(input)
- return input.log()
-
- def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- return grad_output / input
-
-class Exp(Function):
- def forward(ctx, input):
- ret = input.exp()
- ctx.save_for_backward(ret)
- return ret
-
- def backward(ctx, grad_output):
- ret, = ctx.saved_tensors
- return grad_output * ret
-
-# ************* binary ops *************
-
-from tinygrad.ops_cpu import Add, Sub, Mul, unbroadcast
-
-class Pow(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x, y)
- return x ** y
-
- def backward(ctx, grad_output):
- x,y = ctx.saved_tensors
- return unbroadcast(y * (x**(y-1.0)) * grad_output, x.shape), \
- unbroadcast((x**y) * torch.log(x) * grad_output, y.shape)
+from tinygrad.ops_cpu import ReLU, Log, Exp, Add, Sub, Mul, Pow
# ************* reduce ops *************

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -170,6 +170,17 @@ class Tensor:
print(f"warning, {data.shape!r} isn't float32, it's {data.dtype}")
Tensor.did_float_warning = True
+ if device == Device.CPU:
+ # add these functions to ndarray
+ class CPUBuffer(np.ndarray):
+ def log(x):
+ return np.log(x)
+ def exp(x):
+ return np.exp(x)
+ def relu(x):
+ return np.maximum(x, 0)
+ data = data.view(CPUBuffer)
+
if device == Device.GPU:
require_init_gpu()
with ProfileOp("toGPU", [data]):

more similar between torch and numpy