From 3469af7a33fcc62b54448fdc84888db5777aebc1 Mon Sep 17 00:00:00 2001
From: dyKiU <dyKiUgXLfA7U@protonmail.com>
Date: Fri, 21 Aug 2026 15:29:13 +0100
Subject: [PATCH 1/3] fix view assign prior reader state

Keep lazy readers on the value they captured before a view assignment while aliases continue to observe the write. Preserve that state through JIT replay and backward, and avoid realizing disjoint readers.
---
 test/backend/test_assign.py | 166 ++++++++++++++++++++++++++++++++++++
 tinygrad/engine/jit.py      |   5 +-
 tinygrad/tensor.py          | 159 +++++++++++++++++++++++++++++-----
 3 files changed, 306 insertions(+), 24 deletions(-)

diff --git a/test/backend/test_assign.py b/test/backend/test_assign.py
index d10d8888c8536..bec00e195d873 100644
--- a/test/backend/test_assign.py
+++ b/test/backend/test_assign.py
@@ -327,6 +327,172 @@ def test_assign_corealize_order_independent(self):
       self.assertEqual(y.tolist(), [11.0])
       self.assertEqual(x.tolist(), [5.0])
 
+  def test_assign_view_corealize_order_independent(self):
+    for reader_first in (True, False):
+      x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+      reader = x + 10
+      writer = x[:1].assign(Tensor([9], dtype=dtypes.int))
+      Tensor.realize(reader, writer) if reader_first else Tensor.realize(writer, reader)
+      self.assertEqual(reader.tolist(), [11, 12])
+      self.assertEqual(x.tolist(), [9, 2])
+
+  def test_assign_view_reshape_alias_sees_write(self):
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    alias = x.reshape(2, 2)
+    prior = alias + 0
+    x[:1].assign(Tensor([9], dtype=dtypes.int))
+    post = alias + 0
+    Tensor.realize(prior, post)
+    self.assertEqual(prior.tolist(), [[1, 2], [3, 4]])
+    self.assertEqual(post.tolist(), [[9, 2], [3, 4]])
+    self.assertEqual(x.tolist(), [9, 2, 3, 4])
+
+  def test_assign_view_prior_reader_split_realize(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    x[:1].assign(Tensor([9], dtype=dtypes.int)).realize()
+    self.assertEqual(reader.tolist(), [11, 12])
+    self.assertEqual(x.tolist(), [9, 2])
+
+  def test_assign_view_reader_between_assigns_split_realize(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    x[:1].assign(Tensor([9], dtype=dtypes.int))
+    mid = x + 0
+    x[1:].assign(Tensor([8], dtype=dtypes.int))
+    x.realize()
+    self.assertEqual(mid.tolist(), [9, 2])
+    self.assertEqual(x.tolist(), [9, 8])
+
+  def test_assign_view_source_from_prior_reader_split_realize(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    x[:1].assign(reader[1:] * 1).realize()
+    self.assertEqual(x.tolist(), [12, 2])
+    self.assertEqual(reader.tolist(), [11, 12])
+
+  def test_assign_view_source_is_prior_reader(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    x[:].assign(reader).realize()
+    self.assertEqual(x.tolist(), [11, 12])
+    self.assertEqual(reader.tolist(), [11, 12])
+
+  def test_assign_view_prior_reader_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = x + 10
+      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+      return reader.realize()
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11, i+12])
+      self.assertEqual(x.tolist(), [9, i+2])
+
+  def test_assign_view_prior_reader_is_jit_rhs(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = x[:1] + 10
+      x[:1].assign(reader)
+      return reader
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11])
+      self.assertEqual(x.tolist(), [i+11, i+2])
+
+  def test_assign_view_realized_prior_reader_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = (x + 10).realize()
+      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+      return reader
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11, i+12])
+      self.assertEqual(x.tolist(), [9, i+2])
+
+  def test_assign_view_realized_disjoint_reader_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = (x[:1] + 10).realize()
+      x[-1:].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+      return reader
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11])
+      self.assertEqual(x.tolist(), [i+1, 9])
+
+  def test_assign_view_effect_only_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertIsNone(f(x))
+      self.assertEqual(x.tolist(), [9, i+2])
+
+  def test_assign_view_disjoint_reader_stays_lazy(self):
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    reader = x[:1] + 10
+    GlobalCounters.reset()
+    x[-1:].assign(Tensor([9], dtype=dtypes.int))
+    assert_kernel_count(0)
+    self.assertEqual(reader.tolist(), [11])
+    self.assertEqual(x.tolist(), [1, 2, 3, 9])
+
+  def test_assign_view_swap_regions(self):
+    for reverse in (False, True):
+      x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+      a, b = x[:1] * 1, x[1:] * 1
+      w1, w2 = x[:1].assign(b), x[1:].assign(a)
+      Tensor.realize(w2, w1) if reverse else Tensor.realize(w1, w2)
+      self.assertEqual(x.tolist(), [2, 1])
+
+  def test_assign_view_backward_prior_loss(self):
+    w = Tensor([1.0, 2.0]).contiguous().realize()
+    loss = (w * w).sum()
+    w[:1].assign(Tensor([5.0]))
+    loss.backward()
+    self.assertEqual(loss.item(), 5.0)
+    self.assertEqual(w.grad.tolist() if w.grad is not None else None, [2.0, 4.0])
+
+  def test_assign_full_view_backward_prior_loss(self):
+    for realize_first in (False, True):
+      w = Tensor([1.0, 2.0]).contiguous().realize()
+      loss = (w * w).sum()
+      w[:].assign(Tensor([5.0, 6.0]))
+      if realize_first: w.realize()
+      loss.backward()
+      if not realize_first: w.realize()
+      self.assertEqual(loss.item(), 5.0)
+      self.assertEqual(w.grad.tolist() if w.grad is not None else None, [2.0, 4.0])
+
+  def test_assign_view_prior_loss_owner_replaced(self):
+    w = Tensor([1.0, 2.0]).contiguous().realize()
+    loss = (w * w).sum()
+    w[:1].assign(Tensor([5.0]))
+    w.replace(Tensor([100.0, 200.0]).realize())
+    loss.backward()
+    self.assertIsNone(w.grad)
+
+  def test_assign_view_prior_loss_owner_replaced_with_dependency(self):
+    w = Tensor([1.0, 2.0]).contiguous().realize()
+    loss = (w * w).sum()
+    w[:1].assign(Tensor([5.0]))
+    w.replace(w + 100)
+    loss.backward()
+    self.assertIsNone(w.grad)
+
+  def test_assign_bitcast_view_realize_without_touching_base(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    writer = x[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32))
+    writer.realize()
+    self.assertEqual(x.uop.buffer.numpy().tolist(), [4.0, 3.0, 3.0, 4.0])
+
   def test_assign_contiguous(self):
     b = Tensor.arange(16).reshape(4,4).clone().realize()
     a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)
diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py
index bfb7f29d3572f..f23d36d0ffb6f 100644
--- a/tinygrad/engine/jit.py
+++ b/tinygrad/engine/jit.py
@@ -250,7 +250,10 @@ def __call__(self, *args, **kwargs) -> ReturnType:
       try:
         ret = self.fxn(*args, **kwargs)
         if len(params:=get_parameters(ret)): Tensor.realize(*params)
-      finally: capturing.clear()
+        Tensor._flush_capture_effects()
+      finally:
+        capturing.clear()
+        Tensor._clear_capture_effects()
       if not len(self._linears): raise JitError("didn't JIT anything!")
       _check_no_non_tensor_return(ret)
       if DEBUG >= 1: print(f"JIT captured {len(self._linears)} linears with {len(input_buf_uops)} inputs")
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 24c39a1ae35aa..6e08f42a42260 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -241,22 +241,68 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
 
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
-all_tensors: dict[weakref.ref[Tensor], None] = {}
-def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
+def _view_base(u:UOp) -> UOp:
+  while not u.has_buffer_identity() and u.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH}: u = u.src[0]
+  return u
+
+def _views_through(u:UOp, base:UOp) -> bool:
+  # aliases can stack view ops on top of base (e.g. RESHAPE(BUFFER) when base is BUFFER), so walk through buffer-identity nodes too
+  while u is not base and u.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH}: u = u.src[0]
+  return u is base
+
+def _storage_range(u:UOp, base:UOp) -> tuple[UOp, int, int]|None:
+  if not _views_through(u, base) or not all_int(u.shape) or (cv:=u.contiguous_view()) is None: return None
+  root, offset = cv
+  byte_offset = offset * root.element_size()
+  return root, byte_offset, byte_offset + int(u.numel()) * u.element_size()
+
+def _reader_overlaps(nodes:dict[UOp, None], base:UOp, write_range:tuple[UOp, int, int]|None) -> bool:
+  if write_range is None: return True
+  read_views = [s for x in nodes if not _views_through(x, base) for s in x.src if _views_through(s, base)]
+  if not read_views: return True
+  for view in read_views:
+    if (read_range:=_storage_range(view, base)) is None or read_range[0] is not write_range[0]: return True
+    if read_range[1] < write_range[2] and write_range[1] < read_range[2]: return True
+  return False
+
+all_tensors: dict[weakref.ref[Tensor], int] = {}
+_identity_views: dict[weakref.ref[Tensor], None] = {}
+_snapshot_grad_owners: weakref.WeakKeyDictionary[UOp, tuple[tuple[weakref.ref[Tensor], int], ...]] = weakref.WeakKeyDictionary()
+_capture_effects: dict[UOp, Tensor] = {}
+
+def _add_snapshot_grad_owner(target:UOp, owner:weakref.ref[Tensor]) -> None:
+  owners = _snapshot_grad_owners.get(target, ())
+  version = all_tensors[owner]
+  if any(tref() is owner() and owner_version == version for tref,owner_version in owners): return
+  _snapshot_grad_owners[target] = owners+((owner, version),)
+
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, alias_base:UOp|None=None, invert:bool=False,
+                          skip:tuple["Tensor", ...]=()) -> list["Tensor"]:
   with cpu_profile(TracingKey(name), "TINY"):
-    # get tensors in scope
+    # get tensors in scope. with alias_base, only tensors viewing through it (or with invert, only tensors NOT viewing through it)
     in_scope: dict[UOp, bool] = {}
     def visitor(node: UOp) -> bool: return True if node in applied_map else any(in_scope.get(s, False) for s in node.src)
-    scope_tensors: list[Tensor] = [t for tref in list(all_tensors) if (t:=tref()) is not None and t.uop.topovisit(visitor, in_scope)]
+    scope_tensors: list[Tensor] = [t for tref in list(all_tensors) if (t:=tref()) is not None and not any(t is s for s in skip) and
+                                  (alias_base is None or _views_through(t.uop, alias_base) != invert) and t.uop.topovisit(visitor, in_scope)]
 
     # get all Tensors and apply the map. always walk: replace exactly the nodes the map names, values are final
     sink = UOp.sink(*[t.uop for t in scope_tensors])
     new_sink = sink.substitute(applied_map, name=f"substitute {name}", walk=True)
 
     # set the relevant uop to the realized UOps
+    changed: list[Tensor] = []
     for t,s,ns in zip(scope_tensors, sink.src, new_sink.src):
       if s is ns: continue
       t.uop = ns
+      changed.append(t)
+    return changed
+
+def _live_alias_assign_sources(big_sink:UOp) -> tuple[UOp, ...]:
+  """Keep a live assignment source only when it reads the storage being assigned."""
+  store_sources = {u.src[1] for u in big_sink.toposort() if u.op is Ops.STORE and _view_base(u.src[0]) in u.src[1].toposort()}
+  if not store_sources: return ()
+  return tuple(t.uop for tref in list(all_tensors) if (t:=tref()) is not None and t.uop in store_sources and
+               not t.uop.has_buffer_identity() and t.uop not in big_sink.src)
 
 # **** Tensor helper functions ****
 
@@ -323,10 +369,12 @@ def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.P
     if _dtype is not None: self.uop = self.uop.cast(_dtype)
 
     # add to all_tensors after construction succeeds
-    all_tensors[weakref.ref(self)] = None
+    all_tensors[weakref.ref(self)] = 0
 
   @suppress_finalizing
-  def __del__(self): all_tensors.pop(weakref.ref(self), None)
+  def __del__(self):
+    all_tensors.pop(weakref.ref(self), None)
+    _identity_views.pop(weakref.ref(self), None)
 
   def _apply_uop(self, fxn:Callable[..., UOp], *x:Tensor, **kwargs) -> Tensor:
     srcs = (self,)+x
@@ -336,7 +384,12 @@ def _apply_uop(self, fxn:Callable[..., UOp], *x:Tensor, **kwargs) -> Tensor:
     ret = Tensor.__new__(Tensor)
     ret.uop, ret.grad, ret.is_param = new_uop, None, True
     # add to all_tensors after construction succeeds
-    all_tensors[weakref.ref(ret)] = None
+    all_tensors[weakref.ref(ret)] = 0
+    return ret
+
+  def __getitem__(self, indices) -> Tensor:
+    ret = super().__getitem__(indices)
+    if ret is self: _identity_views[weakref.ref(self)] = None
     return ret
 
   # alu, _uop, _wrap_uop and const are used by the mixins
@@ -403,7 +456,9 @@ def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
     # weakness ends where storage begins
     if any(t.dtype in dtypes.weaks and t.uop.device is not None for t in (self,)+lst):
       raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
-    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
+    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
+    big_sink = UOp.sink(*big_sink.src, *_live_alias_assign_sources(big_sink))
+    big_sink, becomes_map = transform_to_call(big_sink)
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
 
@@ -416,17 +471,28 @@ def schedule_linear(self, *lst:Tensor) -> UOp:
   @disable_gc()
   def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Triggers the computation needed to create these Tensor(s)."""
-    to_realize = [x for x in (self,)+lst if not x.uop.is_virtual and not x.uop.has_buffer_identity()]
+    from tinygrad.engine.realize import capturing
+    effects = tuple(_capture_effects.values()) if capturing else ()
+    to_realize = [x for x in (self,)+lst+effects if not x.uop.is_virtual and not x.uop.has_buffer_identity()]
     if len(to_realize):
       run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats)
+    if effects: _capture_effects.clear()
     return self
 
+  @staticmethod
+  def _clear_capture_effects(): _capture_effects.clear()
+
+  @staticmethod
+  def _flush_capture_effects():
+    if _capture_effects: next(iter(_capture_effects.values())).realize()
+
   def replace(self, x:Tensor) -> Tensor:
     """
     Replaces the data of this tensor with the data of another tensor. Only the shape of the tensors must match.
     """
     # used for replacing a Tensor with a new version of it (potentially with a different device and dtype)
     assert self.shape == x.shape, f"replace shape mismatch {self.shape} != {x.shape}"
+    all_tensors[weakref.ref(self)] += 1
     self.uop = x.uop
     return self
 
@@ -450,14 +516,39 @@ def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
       return self
     # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
     assign = self.uop.after(self.uop.store(x.uop))
-    ib = self.uop
-    while not ib.has_buffer_identity() and ib.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH}: ib = ib.src[0]
-    if ib is not self.uop and ib.has_buffer_identity(after_ok=True):
-      # view assign: replace at the buffer-identity level (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
-      _apply_map_to_tensors({ib: ib.after(assign)}, name="Embed View Assign")
+    ib = _view_base(self.uop)
+    self_ref = weakref.ref(self)
+    is_view_assign = ib is not self.uop or self_ref in _identity_views
+    _identity_views.pop(self_ref, None)
+    if is_view_assign and ib.has_buffer_identity(after_ok=True):
+      live_tensors = [(tref, t) for tref in list(all_tensors) if (t:=tref()) is not None]
+      reader_graphs = [(t, nodes) for _,t in live_tensors if not _views_through(t.uop, ib) and ib in (nodes:=t.uop.toposort())]
+      write_range = _storage_range(self.uop, ib)
+      disjoint_readers = tuple(t for t,nodes in reader_graphs if not _reader_overlaps(nodes, ib, write_range))
+      # remember aliases before their public graphs move to the post-assign state. if a differentiable reader uses one,
+      # its corresponding snapshot node remains the gradient target for that Tensor.
+      reader_nodes = {u for _,nodes in reader_graphs for u in nodes}
+      grad_owners = [(tref, t.uop) for tref,t in live_tensors if t.is_floating_point() and t.device is not None and
+                     _views_through(t.uop, ib) and t.uop in reader_nodes]
+      for tref, old_uop in grad_owners: _add_snapshot_grad_owner(old_uop, tref)
+      # prior readers (non-alias tensors reading this buffer) must keep the value they captured no matter when they
+      # realize: repoint overlapping readers at a snapshot copy. disjoint contiguous readers can use the original buffer.
+      overlapping_readers = [t for t,_ in reader_graphs if t is not x and not any(t is d for d in disjoint_readers)]
+      if overlapping_readers:
+        readers = _apply_map_to_tensors({ib: (snap:=ib.clone())}, name="Snapshot Prior Readers", alias_base=ib, invert=True,
+                                        skip=(x,)+disjoint_readers)
+      else: readers = []
+      _apply_map_to_tensors({ib: ib.after(assign)}, name="Embed View Assign", alias_base=ib)
+      # the snapshot graph only references the pre-assign chain, so this can't schedule the new write
+      if readers:
+        realized_snap = Tensor(snap).realize().uop
+        for tref, old_uop in grad_owners:
+          snapshot_uop = old_uop.substitute({ib: realized_snap}, name="snapshot gradient target", walk=True)
+          _add_snapshot_grad_owner(snapshot_uop, tref)
     else:
-      # simple assign
       self.uop = assign
+    from tinygrad.engine.realize import capturing
+    if capturing and ib.has_buffer_identity(after_ok=True): _capture_effects[self.uop.buf_uop] = Tensor(self.uop)
     return self
 
   def _buffer(self) -> Buffer:
@@ -662,15 +753,37 @@ def backward(self, gradient:Tensor|None=None) -> Tensor:
     ```
     """
     all_uops = self.uop.toposort()
-    # backward fills .grad for every in-scope float tensor with a device
-    tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and \
-                                       t.uop in all_uops and t.is_floating_point() and t.device is not None]
+    historical_owners = {target:owners for target in all_uops if (owners:=_snapshot_grad_owners.get(target, ()))}
+    if not historical_owners:
+      tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and
+                                         t.uop in all_uops and t.is_floating_point() and t.device is not None]
+      for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)):
+        assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
+        if g.device is None: g = g.clone(device=t.device)
+        if t.grad is None: t.grad = g
+        else: t.grad.assign(t.grad + g.to(t.grad.device))
+      return self
+    # Map each graph target to the live Tensors that own its gradient. Assign snapshots use a historical graph target
+    # whose gradient still belongs to the Tensor that now exposes the post-assign state.
+    target_owners: dict[UOp, list[Tensor]] = {}
+    def add_owner(target:UOp, owner:Tensor) -> None:
+      if owner.is_floating_point() and owner.device is not None and not any(owner is x for x in target_owners.setdefault(target, [])):
+        target_owners[target].append(owner)
+    for tref in all_tensors:
+      if (t:=tref()) is not None and t.uop in all_uops: add_owner(t.uop, t)
+    for target,owners in historical_owners.items():
+      for tref,owner_version in owners:
+        if (t:=tref()) is not None and all_tensors.get(tref) == owner_version: add_owner(target, t)
+    targets = list(target_owners)
+    gradient_targets = [owners[0] if owners[0].uop is target else Tensor(target) for target,owners in target_owners.items()]
     # clear contexts
-    for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)):
-      assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
-      if g.device is None: g = g.clone(device=t.device)
-      if t.grad is None: t.grad = g
-      else: t.grad.assign(t.grad + g.to(t.grad.device))
+    for target,g in zip(targets, self.gradient(*gradient_targets, gradient=gradient)):
+      for i,t in enumerate(target_owners[target]):
+        owner_grad = g if i == 0 else Tensor(g.uop)
+        assert owner_grad.shape == t.shape, f"grad shape must match tensor shape, {owner_grad.shape!r} != {t.shape!r}"
+        if owner_grad.device is None: owner_grad = owner_grad.clone(device=t.device)
+        if t.grad is None: t.grad = owner_grad
+        else: t.grad.assign(t.grad + owner_grad.to(t.grad.device))
     return self
 
   # ***** movement ops *****

From db5670cfb0633517624b047a1b90120bdea5cec5 Mon Sep 17 00:00:00 2001
From: dyKiU <dyKiUgXLfA7U@protonmail.com>
Date: Fri, 21 Aug 2026 15:30:47 +0100
Subject: [PATCH 2/3] skip view assign ranges without readers

The write range only classifies prior readers, so avoid its graph rewrite when no readers exist. This restores chained view assignment from 73.03 ms to 0.44 ms per assignment while retaining overlap checks when readers are present.
---
 test/backend/test_assign.py | 29 +++++++++++++++++++++++++++++
 tinygrad/tensor.py          |  3 ++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/test/backend/test_assign.py b/test/backend/test_assign.py
index bec00e195d873..f7eab4dd2e035 100644
--- a/test/backend/test_assign.py
+++ b/test/backend/test_assign.py
@@ -444,6 +444,35 @@ def test_assign_view_disjoint_reader_stays_lazy(self):
     self.assertEqual(reader.tolist(), [11])
     self.assertEqual(x.tolist(), [1, 2, 3, 9])
 
+  def _count_storage_range(self, fxn):
+    # _storage_range runs a full graph_rewrite (via UOp.contiguous_view), so it dominates assign cost
+    import tinygrad.tensor as tensor_module
+    calls, orig = [], tensor_module._storage_range
+    def counting(u, base):
+      calls.append(u)
+      return orig(u, base)
+    tensor_module._storage_range = counting
+    try: fxn()
+    finally: tensor_module._storage_range = orig
+    return len(calls)
+
+  def test_assign_view_no_readers_skips_write_range(self):
+    # the write range only decides which prior readers overlap the write. with no prior readers there is
+    # nothing to decide, and computing it anyway made chained view assigns quadratic in chain length
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    n = self._count_storage_range(lambda: x[:1].assign(Tensor([9], dtype=dtypes.int)))
+    self.assertEqual(n, 0, "view assign with no prior readers must not compute the write range")
+    self.assertEqual(x.tolist(), [9, 2, 3, 4])
+
+  def test_assign_view_with_reader_still_checks_overlap(self):
+    # guard the other way: skipping the write range when readers DO exist would lose overlap detection
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    reader = x[:1] + 10
+    n = self._count_storage_range(lambda: x[-1:].assign(Tensor([9], dtype=dtypes.int)))
+    self.assertGreater(n, 0, "with a prior reader the write range is needed to classify overlap")
+    self.assertEqual(reader.tolist(), [11])
+    self.assertEqual(x.tolist(), [1, 2, 3, 9])
+
   def test_assign_view_swap_regions(self):
     for reverse in (False, True):
       x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 6e08f42a42260..b98a9df63838e 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -523,7 +523,8 @@ def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
     if is_view_assign and ib.has_buffer_identity(after_ok=True):
       live_tensors = [(tref, t) for tref in list(all_tensors) if (t:=tref()) is not None]
       reader_graphs = [(t, nodes) for _,t in live_tensors if not _views_through(t.uop, ib) and ib in (nodes:=t.uop.toposort())]
-      write_range = _storage_range(self.uop, ib)
+      # _storage_range runs a graph_rewrite over the whole AFTER chain, so only pay it when a reader needs classifying
+      write_range = _storage_range(self.uop, ib) if reader_graphs else None
       disjoint_readers = tuple(t for t,nodes in reader_graphs if not _reader_overlaps(nodes, ib, write_range))
       # remember aliases before their public graphs move to the post-assign state. if a differentiable reader uses one,
       # its corresponding snapshot node remains the gradient target for that Tensor.

From 0de35ac5a2fdf5f3b84d21af0bcbf2d5521a6e3d Mon Sep 17 00:00:00 2001
From: dyKiU <dyKiUgXLfA7U@protonmail.com>
Date: Sat, 22 Aug 2026 16:21:35 +0100
Subject: [PATCH 3/3] preserve lazy view assign snapshots

---
 test/backend/test_assign.py | 14 ++++++++++++--
 tinygrad/engine/jit.py      |  1 -
 tinygrad/tensor.py          | 30 +++++++++++++++++++++---------
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/test/backend/test_assign.py b/test/backend/test_assign.py
index f7eab4dd2e035..4c93811246e3f 100644
--- a/test/backend/test_assign.py
+++ b/test/backend/test_assign.py
@@ -425,10 +425,10 @@ def f(x:Tensor):
       self.assertEqual(f(x).tolist(), [i+11])
       self.assertEqual(x.tolist(), [i+1, 9])
 
-  def test_assign_view_effect_only_jit_replay(self):
+  def test_assign_view_explicit_effect_jit_replay(self):
     @TinyJit
     def f(x:Tensor):
-      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int)).realize()
 
     for i in range(4):
       x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
@@ -444,6 +444,16 @@ def test_assign_view_disjoint_reader_stays_lazy(self):
     self.assertEqual(reader.tolist(), [11])
     self.assertEqual(x.tolist(), [1, 2, 3, 9])
 
+  def test_assign_view_overlapping_reader_stays_lazy(self):
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    value = Tensor([9], dtype=dtypes.int)
+    GlobalCounters.reset()
+    x[:1].assign(value)
+    assert_kernel_count(0)
+    self.assertEqual(reader.tolist(), [11, 12, 13, 14])
+    self.assertEqual(x.tolist(), [9, 2, 3, 4])
+
   def _count_storage_range(self, fxn):
     # _storage_range runs a full graph_rewrite (via UOp.contiguous_view), so it dominates assign cost
     import tinygrad.tensor as tensor_module
diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py
index f23d36d0ffb6f..3bb1a0cf34ad3 100644
--- a/tinygrad/engine/jit.py
+++ b/tinygrad/engine/jit.py
@@ -250,7 +250,6 @@ def __call__(self, *args, **kwargs) -> ReturnType:
       try:
         ret = self.fxn(*args, **kwargs)
         if len(params:=get_parameters(ret)): Tensor.realize(*params)
-        Tensor._flush_capture_effects()
       finally:
         capturing.clear()
         Tensor._clear_capture_effects()
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index b98a9df63838e..66fba6008eccb 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -268,6 +268,7 @@ def _reader_overlaps(nodes:dict[UOp, None], base:UOp, write_range:tuple[UOp, int
 all_tensors: dict[weakref.ref[Tensor], int] = {}
 _identity_views: dict[weakref.ref[Tensor], None] = {}
 _snapshot_grad_owners: weakref.WeakKeyDictionary[UOp, tuple[tuple[weakref.ref[Tensor], int], ...]] = weakref.WeakKeyDictionary()
+_assign_snapshots: weakref.WeakKeyDictionary[UOp, UOp] = weakref.WeakKeyDictionary()
 _capture_effects: dict[UOp, Tensor] = {}
 
 def _add_snapshot_grad_owner(target:UOp, owner:weakref.ref[Tensor]) -> None:
@@ -276,6 +277,14 @@ def _add_snapshot_grad_owner(target:UOp, owner:weakref.ref[Tensor]) -> None:
   if any(tref() is owner() and owner_version == version for tref,owner_version in owners): return
   _snapshot_grad_owners[target] = owners+((owner, version),)
 
+def _transfer_snapshot_metadata(applied_map:dict[UOp, UOp]) -> None:
+  for target,replacement in applied_map.items():
+    if not (owners:=_snapshot_grad_owners.get(target)): continue
+    existing = _snapshot_grad_owners.get(replacement, ())
+    _snapshot_grad_owners[replacement] = existing+tuple(owner for owner in owners if owner not in existing)
+  for assign,snapshot in list(_assign_snapshots.items()):
+    if (new_snapshot:=applied_map.get(snapshot)) is not None: _assign_snapshots[assign] = new_snapshot
+
 def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, alias_base:UOp|None=None, invert:bool=False,
                           skip:tuple["Tensor", ...]=()) -> list["Tensor"]:
   with cpu_profile(TracingKey(name), "TINY"):
@@ -304,6 +313,9 @@ def _live_alias_assign_sources(big_sink:UOp) -> tuple[UOp, ...]:
   return tuple(t.uop for tref in list(all_tensors) if (t:=tref()) is not None and t.uop in store_sources and
                not t.uop.has_buffer_identity() and t.uop not in big_sink.src)
 
+def _live_assign_snapshots(big_sink:UOp) -> tuple[UOp, ...]:
+  return tuple(snapshot for u in big_sink.toposort() if (snapshot:=_assign_snapshots.get(u)) is not None)
+
 # **** Tensor helper functions ****
 
 def is_numpy_ndarray(x) -> "TypeGuard[numpy.ndarray]": return str(type(x)) == "<class 'numpy.ndarray'>"
@@ -448,7 +460,9 @@ def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None)
   def callify(self, *lst:Tensor) -> Tensor:
     big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
     big_sink, buffer_map = transform_to_call(big_sink)
-    _apply_map_to_tensors({x:y.after(big_sink) for x,y in buffer_map.items()}, name="callify")
+    applied_map = {x:y.after(big_sink) for x,y in buffer_map.items()}
+    _transfer_snapshot_metadata(applied_map)
+    _apply_map_to_tensors(applied_map, name="callify")
     return self
 
   def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
@@ -458,7 +472,10 @@ def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
       raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
     big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
     big_sink = UOp.sink(*big_sink.src, *_live_alias_assign_sources(big_sink))
+    # snapshots are separate roots so WAR scheduling orders their reads before the corresponding writes
+    big_sink = UOp.sink(*big_sink.src, *_live_assign_snapshots(big_sink))
     big_sink, becomes_map = transform_to_call(big_sink)
+    _transfer_snapshot_metadata(becomes_map)
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
 
@@ -482,10 +499,6 @@ def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
   @staticmethod
   def _clear_capture_effects(): _capture_effects.clear()
 
-  @staticmethod
-  def _flush_capture_effects():
-    if _capture_effects: next(iter(_capture_effects.values())).realize()
-
   def replace(self, x:Tensor) -> Tensor:
     """
     Replaces the data of this tensor with the data of another tensor. Only the shape of the tensors must match.
@@ -536,15 +549,14 @@ def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
       # realize: repoint overlapping readers at a snapshot copy. disjoint contiguous readers can use the original buffer.
       overlapping_readers = [t for t,_ in reader_graphs if t is not x and not any(t is d for d in disjoint_readers)]
       if overlapping_readers:
-        readers = _apply_map_to_tensors({ib: (snap:=ib.clone())}, name="Snapshot Prior Readers", alias_base=ib, invert=True,
+        readers = _apply_map_to_tensors({ib: (snap:=ib.detach().clone())}, name="Snapshot Prior Readers", alias_base=ib, invert=True,
                                         skip=(x,)+disjoint_readers)
+        _assign_snapshots[assign] = snap
       else: readers = []
       _apply_map_to_tensors({ib: ib.after(assign)}, name="Embed View Assign", alias_base=ib)
-      # the snapshot graph only references the pre-assign chain, so this can't schedule the new write
       if readers:
-        realized_snap = Tensor(snap).realize().uop
         for tref, old_uop in grad_owners:
-          snapshot_uop = old_uop.substitute({ib: realized_snap}, name="snapshot gradient target", walk=True)
+          snapshot_uop = old_uop.substitute({ib: snap}, name="snapshot gradient target", walk=True)
           _add_snapshot_grad_owner(snapshot_uop, tref)
     else:
       self.uop = assign