diff --git a/tinygrad/engine/lazy.py b/tinygrad/engine/lazy.py index e61ea4e3f5..74bd41a246 100644 --- a/tinygrad/engine/lazy.py +++ b/tinygrad/engine/lazy.py @@ -32,6 +32,7 @@ class LazyBuffer(MathTrait): # properties on base self.op, self.arg, self.srcs = op, arg, srcs # this is a UOp, except the src is LazyBuffers and not UOps assert self.op is not Ops.ASSIGN or srcs[0].base.realized is not None, "assign target must be realized" + assert all_same([x.st.shape for x in self.srcs]), f"src shape mismatch! {self.srcs}" if self.op is Ops.BUFFER_VIEW: # some LazyBuffers can be processed with only a view, no AST required @@ -81,7 +82,7 @@ class LazyBuffer(MathTrait): 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(Ops.ASSIGN, self.shape, self.dtype, self.device, arg=None if self.st.contiguous else self.st, - src=(self.base, x), enable_cache=True) + src=(self, x), enable_cache=True) # NOTE: assign to VIEW is fine def can_view(self): return (self.st.consecutive and not self.is_unrealized_const() and not isinstance(self.dtype, ImageDType) and diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index fd9f5a8fca..78f07cfd26 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -70,7 +70,7 @@ def to_uop(buf:LazyBuffer, ctx:ScheduleContext, buffers:Dict[UOp, Buffer], cache op = None elif buf.op is Ops.ASSIGN: target, new_val = [to_uop(x, ctx, buffers, cache) for x in buf.srcs] - ctx.assigns.add(ubuf:=target.buf_uop) + ctx.assigns.add(ubuf:=target.base.buf_uop) op = UOp(Ops.ASSIGN, dtype.base, (ubuf, new_val), buf.arg) else: ubuf = UOp.new_buffer(buf.device, buf.size, dtype)