DeviceBuffer has abstract methods
MODIFY
test/lib_test_ast.py
test/lib_test_ast.py
import sys
import numpy as np
-from typing import Dict
+from typing import Dict, Type
from tinygrad.ast import ASTKernel
from tinygrad.llops.ops_cpu import CPUBuffer
from tinygrad.ops import DeviceBuffer
@@ -8,15 +8,16 @@ from tinygrad.lazy import realize_buffers
in_test = False
test_cnt = 0
-def test_ast(k:ASTKernel, device:DeviceBuffer=CPUBuffer):
+def test_ast(k:ASTKernel, device:Type[DeviceBuffer]=CPUBuffer):
global in_test, test_cnt
if in_test: return
in_test = True
print("testing AST", test_cnt)
test_cnt += 1
cpubufs : Dict[DeviceBuffer, DeviceBuffer] = {x:device.fromCPU(x.toCPU()) for x in k.bufs}
- real_out = cpubufs[k.bufs[0]]
- test_out = device.exec_ast(realize_buffers(cpubufs, k.ast))
+ real_out = cpubufs[k.bufs[0]].toCPU()
+ assert hasattr(device, 'exec_ast')
+ test_out = device.exec_ast(realize_buffers(cpubufs, k.ast)).toCPU()
if not np.allclose(real_out, test_out, atol=1e-4):
print("MISMATCH")
print(k.print())

MODIFY
tinygrad/ops.py
tinygrad/ops.py
+from __future__ import annotations
import os
+import numpy as np
from enum import Enum
from typing import Union, Type, NamedTuple, Tuple, Any, List
import functools, operator
@@ -32,10 +34,15 @@ def get_lazyops(op:LazyOp) -> List[LazyOp]: return functools.reduce(operator.add
# a placeholder class to extend by the exec classes
class DeviceBuffer:
shape: Any # should be Tuple[int, ...] but ndarray and torch.tensor have incompatible types
+ @staticmethod
+ def fromCPU(x:np.ndarray) -> DeviceBuffer: raise NotImplementedError("must be implemented")
+ def toCPU(self:DeviceBuffer) -> np.ndarray: raise NotImplementedError("must be implemented")
+ @classmethod
+ def exec_ast(cls, ast:LazyOp): raise NotImplementedError("must be implemented")
# extend this if you don't have an exec_ast function
# used in CPUBuffer and TorchBuffer
-class GenericExecAST(DeviceBuffer):
+class GenericExecAST(DeviceBuffer): # pylint: disable=abstract-method
@classmethod
def exec_ast(cls, ast:LazyOp, preprocess=lambda x: x):
srcs = [cls.exec_ast(x, preprocess) if isinstance(x, LazyOp) else preprocess(x) for x in ast.src]
@@ -58,7 +65,7 @@ class GenericExecAST(DeviceBuffer):
class GlobalCounters:
global_ops, global_mem = 0, 0
-class GenericShape(GenericExecAST):
+class GenericShape(GenericExecAST): # pylint: disable=abstract-method
def __init__(self, shape, flops=0): self.shape, self.flops = shape, flops
def unary_op(self, op:UnaryOps): return GenericShape(self.shape, self.flops + prod(self.shape))
def binary_op(self, op:BinaryOps, y): return GenericShape(self.shape, self.flops + y.flops + prod(self.shape))
@@ -70,14 +77,11 @@ def get_lazyop_info(ast:LazyOp): return GenericShape.exec_ast(ast, lambda x: Gen
# assumes you are using ShapeTracker
# used in GPUBuffer and LLVMBuffer
-class ExplicitExecAST(DeviceBuffer):
+class ExplicitExecAST(DeviceBuffer): # pylint: disable=abstract-method
def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf=None):
self.st = shape if isinstance(shape, ShapeTracker) else ShapeTracker(tuple(shape))
self.shape = self.st.shape
- @classmethod
- def exec_ast(cls, ast:LazyOp): raise NotImplementedError("must be implemented")
-
# universal
def unary_op(self, op:UnaryOps): return type(self)(self.shape).exec_ast(LazyOp(op=op, src=(self,)))
def binary_op(self, op:BinaryOps, y): return type(self)(self.shape).exec_ast(LazyOp(op=op, src=(self, y)))

DeviceBuffer has abstract methods