From 9fe596ce6e86ea88da01012f290b48ccd750b6f1 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Mon, 4 Nov 2024 14:04:39 +0200
Subject: [PATCH] early assert assign [pr] (#7526)

* early assert assign [pr]

* self

* don't need base
---
 tinygrad/engine/fuse.py | 5 +----
 tinygrad/engine/lazy.py | 1 +
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/tinygrad/engine/fuse.py b/tinygrad/engine/fuse.py
index 312ca7fbc7..cbf6236f2a 100644
--- a/tinygrad/engine/fuse.py
+++ b/tinygrad/engine/fuse.py
@@ -34,10 +34,7 @@ def _recurse_lb(buf:LazyBuffer, realizes:Dict[LazyBuffer, None], allbufs:Dict[La
   if ctx.buf_uops[buf.buffer] in ctx.realizes: realizes[buf] = None
   if buf.op in GroupOp.Reduce and buf.srcs[0].base.op is buf.op and buf.srcs[0] is not buf.srcs[0].base: double_reduces[buf] = None
   allbufs[buf] = None
-  if buf.op is MetaOps.ASSIGN:
-    assign_targets[(target:=buf.srcs[0])] = buf
-    assert target._base is None, f"assign must be to base {target}"
-    assert target.is_realized(), f"assign must be already realized to schedule {target}"
+  if buf.op is MetaOps.ASSIGN: assign_targets[buf.srcs[0]] = buf
   for x in buf.srcs:
     if x.base.realized is None: children[x.base][buf] = None
     _recurse_lb(x, realizes, allbufs, simple_pads, children, assign_targets, double_reduces, ctx)
diff --git a/tinygrad/engine/lazy.py b/tinygrad/engine/lazy.py
index 30b3b78e8c..438d71014b 100644
--- a/tinygrad/engine/lazy.py
+++ b/tinygrad/engine/lazy.py
@@ -80,6 +80,7 @@ class LazyBuffer(MathTrait):
 
   def assign(self, x:LazyBuffer) -> LazyBuffer:
     assert x.size == self.size, f"assign target must have same size {self.size=} != {x.size=}"
+    assert self.is_realized(), f"assign target must be realized {self}"
     return LazyBuffer.metaop(MetaOps.ASSIGN, self.shape, self.dtype, self.device, arg=() if self.st.contiguous else (self.st,),
                              src=(self.base, x), enable_cache=True)