diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index d809beae99..f9137cd7fb 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -111,7 +111,7 @@ pm_post_sched_cache = PatternMatcher([ ]) def _buffer_like(x:UOp): - buffer = UOp.new_buffer(x.device, x.shard_size, x.dtype).reshape(x.max_shard_shape).shrink_to(x.shard_shape) + buffer = UOp.new_buffer(x.device, x.shard_size, x.dtype).reshape(x.max_shard_shape) if isinstance(x.device, tuple) and x.axis is not None: buffer = buffer.multi(x.axis) return buffer @@ -128,12 +128,14 @@ def contig_to_assign(ctx:dict[UOp,UOp|None], x:UOp): # we can't realize any disk tensors if x in ctx: del ctx[x] # all copies are assumed to have finished and are now just the buffer - if x.op is Ops.COPY: ctx[x] = _buffer_like(x) + if x.op is Ops.COPY: ctx[x] = _buffer_like(x).shrink_to(x.shard_shape) return None # for contiguous or in buffer_map explicitly if not (x.op is Ops.CONTIGUOUS or (x in ctx and ctx[x] is None)): return None # not sure why the ctx isn't enough, but tag fixes it - ctx[x] = buffer = _buffer_like(x) + buffer = _buffer_like(x) + # buffer_map value needs symbolic shrink to preserve tensor shape, but assign target uses max shape (rangeify handles symbolic indexing) + ctx[x] = buffer.shrink_to(x.shard_shape) return buffer.assign(x.src[0] if x.op is Ops.CONTIGUOUS else x.rtag()) pm_build_buffer_map = PatternMatcher([ (UPat(GroupOp.All, name="x"), contig_to_assign), ])