From eaa8d343d86db7f43a87f432810c03ac5961150e Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Sun, 24 Sep 2023 19:03:22 -0400
Subject: [PATCH] Remove str type from map_buffers (#1912)

---
 tinygrad/lazy.py | 6 +++---
 tinygrad/ops.py  | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py
index 10762cf98a..bdf7d070c3 100644
--- a/tinygrad/lazy.py
+++ b/tinygrad/lazy.py
@@ -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
diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index 93d46fe881..e29b9ff152 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -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):