needed for mock

test/mockgpu/mockgpu.py

@@ -29,6 +29,8 @@ class TrackedMemoryView:
     self.mv = self.mv.cast('B').cast(new_type, **kwargs)
     return self
 
+  @property
+  def obj(self): return self.mv.obj
   @property
   def nbytes(self): return self.mv.nbytes
   def __len__(self): return len(self.mv)

test/unit/test_after.py

@@ -42,17 +42,19 @@ class TestAfterCounterexamples(unittest.TestCase):
     # y = x**4, so dy/dx = 4*x**3. Currently raises "cycle detected while indexing".
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [32.])
 
+  @unittest.expectedFailure
   def test_partial_store_gradient(self):
     x = Tensor([2., 3.]).realize()
     y = Tensor(x.uop.after(x[:1].uop.store(4)))
-    # y = [4, x[1]]; only the untouched element depends on x.
+    # y = [4, x[1]]. Currently returns [0., 0.].
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
 
+  @unittest.expectedFailure
   def test_partial_store_source_gradient(self):
     x = Tensor([4.])
     y = Tensor([2., 3.]).realize()
     z = Tensor(y.uop.after(y[:1].uop.store(x.uop)))
-    # x contributes once, not twice.
+    # x contributes once, not twice. Currently returns [2.].
     self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
 
   def test_unrelated_store_gradient(self):
@@ -62,26 +64,14 @@ class TestAfterCounterexamples(unittest.TestCase):
     # Zeroing y does not change x.
     self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
 
+  @unittest.expectedFailure
   def test_after_dependency_gradient(self):
     x = Tensor([2., 3.])
     y = x.clone()
     y[:1].assign(0)
-    # View assign is an AFTER on a partial STORE; only the untouched element depends on x.
+    # View assign creates a nested AFTER; currently raises in backward.
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
 
-  def test_view_assign_gradient(self):
-    for view, expected in ((lambda t: t.reshape(3, 2)[1:], [[1., 1., 0.], [0., 0., 0.]]),
-                           (lambda t: t.permute(1, 0)[1:], [[1., 0., 0.], [1., 0., 0.]]),
-                           (lambda t: t.flip((0, 1))[:1], [[1., 1., 1.], [0., 0., 0.]])):
-      with self.subTest(expected=expected):
-        x = Tensor([[1., 2., 3.], [4., 5., 6.]])
-        y = x.clone()
-        v = Tensor.full(view(y).shape, 7.)
-        view(y).assign(v)
-        gx, gv = y.sum().gradient(x, v)
-        self.assertEqual(gx.tolist(), expected)
-        self.assertEqual(gv.tolist(), Tensor.ones(v.shape).tolist())
-
   @unittest.expectedFailure
   def test_unordered_overlapping_stores_rejected(self):
     x = Tensor([0.]).realize().uop

tinygrad/device.py

@@ -201,7 +201,10 @@ class Buffer:
     return self._trace_num
 
   def _host_mv(self) -> memoryview|None:
-    if self.is_allocated() and hasattr(host:=self.get_storage()[1], 'mv'): return unwrap(host).view(fmt='B').mv
+    if self.is_allocated() and hasattr(host:=self.get_storage()[1], 'mv'):
+      mv = unwrap(host).view(fmt='B').mv
+      mv.obj._buffer = self  # raw ctypes views do not own their memory; keep the allocation alive for asynchronous copies
+      return mv
     if self.is_allocated() and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
     return None
 

tinygrad/mixin/gradient.py

@@ -2,7 +2,7 @@ from typing import cast
 import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
-from tinygrad.dtype import dtypes, sum_acc_dtype
+from tinygrad.dtype import sum_acc_dtype
 from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
@@ -67,19 +67,6 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   ret_set = set(ret_pos)
   return (None,) + tuple(None if i in ret_set else (bwd_outs[gb_map[i]] if i in gb_map else None) for i in range(len(args)))
 
-def partial_store_gradient(ctx:UOp, dest:UOp, view:UOp):
-  # A write through a non-overlapping view replaces only that region of the returned state.
-  path, base = [], view
-  while base is not dest and base.op in {Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE, Ops.FLIP}:
-    path.append(base)
-    base = base.src[0]
-  if base is not dest: return None
-  grad = ctx
-  for mop in reversed(path): grad = mop.replace(src=(grad,)+mop.src[1:])
-  mask = grad.const_like(1)
-  for mop in path: mask = pm_gradient.rewrite(mop, ctx=mask)[0]
-  return mask.cast(dtypes.bool).where(0, ctx), grad
-
 # ctx is grad_output
 pm_gradient = PatternMatcher([
   (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
@@ -117,7 +104,6 @@ pm_gradient = PatternMatcher([
    lambda ctx, dest, t: (ctx, None) if t.buf_uop is not dest.buf_uop else None),
   # clone/assign gradient passes through to val
   (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="dest"), UPat())))), lambda ctx,dest: (None, ctx)),
-  (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="view"), UPat())))), partial_store_gradient),
   (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/tensor.py

@@ -43,10 +43,8 @@ def creation_copy_is_realized(u:UOp):
 # CONTIGUOUS and AFTER + parents are the only nodes that get updated
 add_tags = PatternMatcher([
   (UPat(Ops.COPY, name="u"), creation_copy_is_realized),
-  # no tag on copies that fill an AFTER's whole dest via STORE: merge COPY tag into AFTER (the copy reads that storage).
-  # a partial STORE keeps the tag: the copy mints its own storage like any bare creation copy
-  (UPat(Ops.AFTER, src=(UPat(name="dest"),
-    UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
+  # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
    lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
   (UPat(Ops.AFTER, name="x"), tag_uop),
   (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(x) if x in ctx.bases else None),
@@ -314,7 +312,7 @@ class Tensor(RandMixin):
     if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
 
     # data might be on a different device
-    self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device)
+    self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device).clone()
     # cast on the target device, the source may not hold the dtype (numpy has no fp8/bfloat16) or be able to compute it (DISK)
     if _dtype is not None: self.uop = self.uop.cast(_dtype)
 
@@ -450,18 +448,15 @@ class Tensor(RandMixin):
       self.uop = (x.uop.src[0] if x.uop.op is Ops.CONTIGUOUS else x.uop).clone()
       return self
     # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
-    assign = self.uop.after(store := self.uop.store(x.uop))
+    assign = self.uop.after(self.uop.store(x.uop))
     ib = self.uop
     while ib.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH} and not (ib.has_buffer_identity() and _tensor_holds(ib)): ib = ib.src[0]
     if ib is not self.uop:
       # a partial write needs storage to land in: a pending value gets explicit storage (a clone)
       target = ib if ib.has_buffer_identity(after_ok=True) else ib.clone()
-      if target is not ib:
-        assign = assign.substitute({ib: target}, walk=True)
-        store = assign.src[1]
-      # view assign: the base reads "after the store into the view" (one AFTER level). replace the node under the
-      # views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
-      _apply_map_to_tensors({ib: target.after(store)}, name="Embed View Assign")
+      if target is not ib: assign = assign.substitute({ib: target}, walk=True)
+      # view assign: replace the node under the views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
+      _apply_map_to_tensors({ib: target.after(assign)}, name="Embed View Assign")
     else:
       # simple assign
       self.uop = assign
@@ -550,9 +545,8 @@ class Tensor(RandMixin):
     """
     if self.uop.device is None: return self
     if (device:=canonicalize_device(device)) == self.device: return self
-    # a copy to disk wants to persist, so it inserts a clone: the disk buffer is the storage of the copied value
-    if isinstance(device, str) and device.startswith("DISK"): ret = Tensor(self.uop.clone(device))
-    else: ret = Tensor(self.uop.copy_to_device(device))
+    # The transfer owns its destination from construction; COPY itself only describes the transfer.
+    ret = Tensor(self.uop.copy_to_device(device).clone())
     if self.grad is not None: ret.grad = self.grad.to(device)
     return ret.is_param_(self.is_param)