move ops to ops.py
MODIFY
test/test_llops.py
test/test_llops.py
@@ -4,7 +4,7 @@ import unittest
from tqdm import trange
import numpy as np
from tinygrad.tensor import Device
-from tinygrad.helpers import UnaryOps, BinaryOps, ReduceOps
+from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps
if Device.DEFAULT == Device.GPU:
from tinygrad.llops.ops_gpu import GPUBuffer, sync, unary_op, binary_op, reduce_op

MODIFY
test/test_onnx.py
test/test_onnx.py
@@ -64,6 +64,7 @@ def run_onnx(onnx_model, inputs={}, debug=False):
elif n.op_type == "Sigmoid": ret = inp[0].sigmoid()
elif n.op_type == "Tanh": ret = inp[0].tanh()
elif n.op_type == "Softmax": ret = inp[0].softmax()
+ elif n.op_type == "MatMul": ret = inp[0].matmul(inp[1])
# one liners
elif n.op_type == "Elu": ret = inp[0].elu(alpha=opt['alpha'])
elif n.op_type == "Clip": ret = inp[0].clip(*(inp[1:] if len(inp) > 1 else (opt['min'], opt['max'])))
@@ -73,7 +74,6 @@ def run_onnx(onnx_model, inputs={}, debug=False):
elif n.op_type == "Squeeze": ret = inp[0].reshape([s for i,s in enumerate(inp[0].shape) if i not in opt['axes']])
elif n.op_type == "GlobalAveragePool": ret = inp[0].mean(axis=tuple(range(2, len(inp[0].shape))), keepdim=True)
elif n.op_type == "BatchNormalization": ret = batch_normalize(inp[0], inp[1], inp[2], inp[3], inp[4], opt.get('epsilon', 1e-5))
- elif n.op_type == "MatMul": ret = inp[0].dot(inp[1])
elif n.op_type == "Gemm": ret = inp[0].linear(inp[1].transpose() if opt.get('transB', 0) == 1 else inp[1], inp[2])
elif n.op_type == "Conv":
x,w,b = inp if len(inp) == 3 else (inp[0], inp[1], None)

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -35,11 +35,3 @@ def get_conv_args(x_shape, w_shape, stride, groups):
assert cout % groups == 0
rcout = cout//groups
return conv_args(H, W, groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs)
-
-# TODO: move to ops.py
-from enum import Enum
-UnaryOps = Enum("UnaryOps", ["RELU", "EXP", "LOG", "NEG", "SIGN"])
-BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "A", "CMPEQ"])
-ReduceOps = Enum("ReduceOps", ["SUM", "MAX"])
-MovementOps = Enum("MovementOps", ["RESHAPE", "PERMUTE", "SLICE"])
-ProcessingOps = Enum("ProcessingOps", ["CONV", "CONVT", "CONVDW"])

MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
import numpy as np
-from tinygrad.helpers import get_conv_args, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
+from tinygrad.helpers import get_conv_args
+from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
class CPUBuffer(np.ndarray):
def relu(x): return np.maximum(x, 0)

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
import functools
import numpy as np
import pyopencl as cl
-from tinygrad.helpers import prod, binary_broadcast, get_conv_args, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
+from tinygrad.helpers import prod, binary_broadcast, get_conv_args
+from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
cl_ctx, cl_queue = None, None
def require_init_gpu():

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
@@ -23,7 +23,8 @@ from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, movement_op
# ************* processing ops *************
-from tinygrad.helpers import get_conv_args, ProcessingOps
+from tinygrad.helpers import get_conv_args
+from tinygrad.ops import ProcessingOps
def convdw(x,grad_output,dw,stride,groups):
# NOTE: torch.nn.grad.conv2d_weight is wrong for groups in pytorch, wonder who it affects

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
import numpy as np # TODO: remove this, it's used for np.prod and np.argsort
-from tinygrad.helpers import binary_broadcast, get_conv_args, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
+from tinygrad.helpers import prod, binary_broadcast, get_conv_args
+from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
from tinygrad.tensor import Function
# ************* unary ops *************
@@ -125,7 +126,7 @@ class Pow(Function):
class Reshape(Function):
def forward(ctx, x, shape):
ctx.save_for_backward(x.shape)
- shape = tuple(-np.prod(x.shape) // np.prod(shape) if s == -1 else s for s in shape)
+ shape = tuple(-prod(x.shape) // prod(shape) if s == -1 else s for s in shape)
return ctx.movement_op(MovementOps.RESHAPE, x, ctx.buffer(shape))
def backward(ctx, grad_output):

MODIFY
tinygrad/ops.py
tinygrad/ops.py
# TODO: move Device to here and proxy buffer call
+from enum import Enum
+UnaryOps = Enum("UnaryOps", ["RELU", "EXP", "LOG", "NEG", "SIGN"])
+BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "A", "CMPEQ"])
+ReduceOps = Enum("ReduceOps", ["SUM", "MAX"])
+MovementOps = Enum("MovementOps", ["RESHAPE", "PERMUTE", "SLICE"])
+ProcessingOps = Enum("ProcessingOps", ["CONV", "CONVT", "CONVDW"])
+
import os
DEBUG = int(os.getenv("PRINT_LLOPS", "0"))
class Ops:

move ops to ops.py