hasattr and DeviceBuffer type fixups
MODIFY
test/test_cl_tiler.py
test/test_cl_tiler.py
@@ -27,7 +27,7 @@ def get_pts(*boundary_coords):
c = sum(([i - 1, i, i + 1] for i in boundary_coords), start=[CL().cl_ctx.devices[0].image2d_max_width])
return [[i, j] for i in c for j in c]
-@unittest.skipUnless(getattr(Device, "OPENCL", None) is not None, "Test requires OpenCL")
+@unittest.skipUnless(hasattr(Device, "OPENCL"), "Test requires OpenCL")
class TestCLTiler(unittest.TestCase):
"""Test for CLImage tiling logic, which allows large tensors to fit within limited-size OpenCL images."""

MODIFY
tinygrad/graph.py
tinygrad/graph.py
@@ -40,13 +40,13 @@ def log_op(optype : OpType, op : List[Op], ret : DeviceBuffer, inp : List[Device
if GRAPH:
def nm(x):
global global_num_max
- if getattr(x, 'global_num', None) is None:
+ if not hasattr(x, 'global_num'):
setattr(x, 'global_num', global_num_max)
global_num_max += 1
return f"<<< {x.global_num} >>>"
top_colors = {LoadOps: '#FFFF80', UnaryOps: "#c0c0c0", ReduceOps: "#8080ff", BinaryOps: "#c0c0c0", MovementOps: "#80ff80", ProcessingOps: "#ff8080"}
- dashed = (optype == LoadOps and getattr(ret, "_backing", None) is not None) or (getattr(ret, "st", None) is not None and not ret.st.contiguous)
+ dashed = (optype == LoadOps and hasattr(ret, "_backing")) or (hasattr(ret, "st") and not ret.st.contiguous) # type: ignore
for x in inp:
if len(op) <= 2:

MODIFY
tinygrad/lazy.py
tinygrad/lazy.py
@@ -64,7 +64,7 @@ def _realize_processingops(self:LazyBuffer) -> Tuple[DeviceBuffer, List[DeviceBu
return real_src_x.processing_op(self.op.op, real_src_w, self.op.arg), [real_src_x, real_src_w], ProcessingOps
def _realize_binaryops(self:LazyBuffer) -> Tuple[DeviceBuffer, List[DeviceBuffer], OpType]:
- real_srcs : Dict[LazyBuffer, DeviceBuffer] = {x:None for x in get_buffers(self.op)}
+ real_srcs : Dict[LazyBuffer, Union[None, LazyOp, DeviceBuffer]] = {x:None for x in get_buffers(self.op)}
op_type : OpType = BinaryOps
psrcs : List[Tuple[LazyBuffer, LazyBuffer]] = [(k,x) for k,x in zip(real_srcs.keys(), map(get_movementroot_contiguous, real_srcs.keys())) if x.optype in [ProcessingOps,ReduceOps] and x.realized is None and len(x.children) <= 1 and len(k.children) <= 1]
if len(psrcs) == 1 and MERGE_ONE_REDUCE_INTO_ELEMENTWISE and (self.device != "OPENCL" or self.shape[-1] == 4):
@@ -88,7 +88,7 @@ def _realize_binaryops(self:LazyBuffer) -> Tuple[DeviceBuffer, List[DeviceBuffer
real_srcs[x] = x.realize(self.device)
ret = self.dbuffer.exec_ast(realize_buffers(real_srcs, self.op))
assert ret.shape == self.shape, f"shape mismatch {ret.shape} != {self.shape}"
- return ret, [x for x in real_srcs.values() if not isinstance(x, LazyOp)], op_type
+ return ret, [x for x in real_srcs.values() if not isinstance(x, LazyOp) and x is not None], op_type
_realize = {LoadOps:_realize_loadops, ReduceOps:_realize_reduceops, MovementOps:_realize_movementops, BinaryOps:_realize_binaryops, ProcessingOps:_realize_processingops}
@@ -113,7 +113,7 @@ class LazyBuffer:
return LazyBuffer.lazycache[wop]
def __init__(self, device, shape:Union[ShapeTracker, Tuple[int, ...]], optype:OpType, op:LazyOp):
- if getattr(self, 'device', None) is not None:
+ if hasattr(self, 'device'):
return # cache hit, we return and don't reinit
self.st = shape if isinstance(shape, ShapeTracker) else ShapeTracker(tuple(shape))
self.shape, self.optype, self.op = self.st.shape, optype, op

MODIFY
tinygrad/nn/__init__.py
tinygrad/nn/__init__.py
@@ -34,7 +34,7 @@ class BatchNorm2D:
else:
batch_mean, batch_var = self.running_mean, self.running_var
# NOTE: this can be precomputed for static inference. if you manually update running_var, you have to reset this
- if getattr(self, "batch_invstd", None) is None:
+ if not hasattr(self, "batch_invstd"):
self.batch_invstd = batch_var.add(self.eps)**-0.5
batch_invstd = self.batch_invstd

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -29,14 +29,12 @@ def get_buffers(op:LazyOp) -> List[Any]: return functools.reduce(operator.add, [
def get_lazyops(op:LazyOp) -> List[LazyOp]: return functools.reduce(operator.add, [get_lazyops(x) for x in op.src if isinstance(x, LazyOp)], [op])
# a placeholder class to extend by the exec classes
-#class DeviceBuffer: pass
-
-# TODO: fix the types for this
-DeviceBuffer = Any
+class DeviceBuffer:
+ shape: Any # should be Tuple[int, ...] but ndarray and torch.tensor have imcompatible types
# extend this if you don't have an exec_ast function
# used in CPUBuffer and TorchBuffer
-class GenericExecAST:
+class GenericExecAST(DeviceBuffer):
@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]
@@ -67,8 +65,7 @@ def get_lazyop_shape(ast:LazyOp): return GenericShape.exec_ast(ast, GenericShape
# assumes you are using ShapeTracker
# used in GPUBuffer, OpenCLBuffer, and LLVMBuffer
-# type: ignore
-class ExplicitExecAST:
+class ExplicitExecAST(DeviceBuffer):
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

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -343,7 +343,7 @@ for device in [device for device in Device.__dict__.keys() if device[0] != "_"]:
# register all the mlops "math" operations
def register(name:str, fxn:Function):
- setattr(Tensor, "_"+name if (getattr(Tensor, name, None) is not None) else name, functools.partialmethod(fxn.apply))
+ setattr(Tensor, "_"+name if hasattr(Tensor, name) else name, functools.partialmethod(fxn.apply))
for name, cls in inspect.getmembers(importlib.import_module('tinygrad.mlops'), inspect.isclass):
if name[0] != "_" and name != "Function" and not name.endswith("Ops"):
register(name.lower(), cls)

hasattr and DeviceBuffer type fixups