tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 10:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove str type from map_buffers (#1912)

Browse Source
This commit is contained in:
chenyu
2023-09-25 07:03:22 +08:00
committed by GitHub
parent ae9529e678
commit eaa8d343d8
2 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/lazy.py
+3 -3
View File
@@ -59,8 +59,8 @@ def _ast_binaryops(self:LazyBuffer) -> LazyOp:
# NOTE: these RESHAPEs will return self if they don't change the shape
for x in real_srcs.keys():
if real_srcs[x] is None: real_srcs[x] = x.reshape(intermediate_shape)
# NOTE: cast the type to remove the Optional, and add str to the value Union to match argument types
ast = self.op.map_buffers(cast(Dict[LazyBuffer, Union[LazyOp, LazyBuffer, str]], real_srcs))
# NOTE: cast the type to remove the Optional
ast = self.op.map_buffers(cast(Dict[LazyBuffer, Union[LazyOp, LazyBuffer]], real_srcs))
return LazyOp(MovementOps.RESHAPE, (ast, ), self.shape) if intermediate_shape != self.shape else ast
def _replace_loadops(op:LazyOp) -> Tuple[LazyOp, List[LazyBuffer]]:
@@ -302,7 +302,7 @@ class LazyBuffer:
@property
def buffers(self) -> Tuple[LazyBuffer, ...]: return (self,)
def map_buffers(self, real_srcs: Mapping[LazyBuffer, Union[LazyBuffer, LazyOp, str]]): return real_srcs.get(self, self)
def map_buffers(self, real_srcs: Mapping[LazyBuffer, Union[LazyBuffer, LazyOp]]): return real_srcs.get(self, self)
def get_lazyops(self) -> List[LazyOp]: return []
def replace_with_movement_ops(self: LazyBuffer, ops:List[Tuple[MovementOps, Any]]) -> LazyBuffer:
y = self
tinygrad/ops.py
+2 -2
View File
@@ -51,7 +51,7 @@ class LazyOp:
@property
def key(self): return (self.op, tuple(map(lambda x: getattr(x, "key", x), self.src)), getattr(self.arg, "key", self.arg))
def map_buffers(self, real_srcs: Mapping[LazyBuffer, Union[LazyBuffer, LazyOp, str]]) -> LazyOp: return LazyOp(self.op, tuple([y.map_buffers(real_srcs) for y in self.src]), self.arg)
def map_buffers(self, real_srcs: Mapping[LazyBuffer, Union[LazyBuffer, LazyOp]]) -> LazyOp: return LazyOp(self.op, tuple([y.map_buffers(real_srcs) for y in self.src]), self.arg)
def get_lazyops(self) -> List[LazyOp]: return [self] + [item for x in self.src for item in x.get_lazyops()]
def replace_with_movement_ops(self:LazyOp, ops:List[Tuple[MovementOps, Tuple[Any, ...]]]) -> 'LazyBuffer':
@@ -179,7 +179,7 @@ from tinygrad.shape.symbolic import Variable, sym_infer
class BasicBatchExecutor:
def __init__(self, jit_cache:List[Tuple[Any, Any, Any]]): pass
def exec(self, jit_cache: List[Tuple[Any, Any, Any]], updatable_entries):
for prg, pargs, variables in jit_cache: prg(pargs, variables, jit=True) # type: ignore
for prg, pargs, variables in jit_cache: prg(pargs, variables, jit=True)
class ASTRunner:
def __init__(self, name, prg, global_size:Optional[List[int]]=None, local_size:Optional[List[int]]=None, op_estimate=0, mem_estimate=0, display_name:Optional[str]=None, runtime_args:Optional[dict]=None):