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)