From c85737c200dc4fdf28cda3846a8e00e60eb8a6c3 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Tue, 14 Jan 2025 13:26:56 -0800
Subject: [PATCH] assert to prepare for grad uop [pr] (#8280)

* assert to prepare for grad uop [pr]

* fix test_nn

* fix most of test_tensor

* few more tests

* fix multi

* uniform gradient

* acc_dtype

* any for multi

* fix typing

* fix assert, CAST_BEFORE_VIEW is still the issue

* explict test for CAST_BEFORE_VIEW

---------

Co-authored-by: qazal <77887910+Qazalin@users.noreply.github.com>
---
 test/test_dtype.py         |  3 ++-
 test/test_tensor.py        |  2 +-
 test/unit/test_gradient.py | 12 ++++++++++++
 tinygrad/multi.py          |  2 ++
 tinygrad/tensor.py         |  3 +++
 5 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/test/test_dtype.py b/test/test_dtype.py
index a09ae0c7f9..7a1677bdd0 100644
--- a/test/test_dtype.py
+++ b/test/test_dtype.py
@@ -781,7 +781,8 @@ class TestAutoCastType(unittest.TestCase):
         if DEBUG >= 2:
           print(f"testing {default_dtype=}, {dtype=}")
         a = Tensor([1, 2, 3], dtype=dtype, requires_grad=True)
-        b = (a * 5).sum()
+        # NOTE: this is broken without default_dtype because of CAST_BEFORE_VIEW
+        b = (a * 5).sum(acc_dtype=default_dtype)
         b.backward()  # if there is dtype mismatch, lazy should assert
         assert a.grad.dtype == a.dtype
         np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5])
diff --git a/test/test_tensor.py b/test/test_tensor.py
index a5152caa95..82f4065593 100644
--- a/test/test_tensor.py
+++ b/test/test_tensor.py
@@ -464,7 +464,7 @@ class TestTinygrad(unittest.TestCase):
   def test_repr_with_grad(self):
     a = Tensor([1], requires_grad=True)
     b = Tensor([1])
-    c = (a + b).mean().backward()
+    c = (a + b).sum().backward()
     print(a)
     print(c)
 
diff --git a/test/unit/test_gradient.py b/test/unit/test_gradient.py
index 21182874f2..b36b81f243 100644
--- a/test/unit/test_gradient.py
+++ b/test/unit/test_gradient.py
@@ -93,6 +93,12 @@ class TestTensorGradient(unittest.TestCase):
     dx = z.gradient(x, gradient=dz)[0]
     self.assertListEqual(dx.tolist(), [2.0, 4.0, 6.0])
 
+  def test_cast_before_view(self):
+    x = Tensor([1.0, 1, 1, 1])
+    x_reshaped = x.reshape(2,2)
+    x_casted = x_reshaped.cast(dtypes.float16)
+    x_casted.mean().gradient(x_reshaped)
+
 class TestRealizeMeansRealize(unittest.TestCase):
   def test_randn_realizes(self):
     x = Tensor.randn(2, 3, 64, 64, requires_grad=True).realize()
@@ -104,5 +110,11 @@ class TestRealizeMeansRealize(unittest.TestCase):
     print(x.lazydata)
     self.assertEqual(x.lazydata.op, Ops.VIEW)
 
+  # NOTE: even though it doesn't realize, this seems fine
+  def test_uniform_gradient(self):
+    x = Tensor.uniform(16, 3, 3, 3, requires_grad=True).realize()
+    y = x * 2
+    y.sum().gradient(x)[0].realize()
+
 if __name__ == '__main__':
   unittest.main()
diff --git a/tinygrad/multi.py b/tinygrad/multi.py
index 919d33a97f..ab521e9978 100644
--- a/tinygrad/multi.py
+++ b/tinygrad/multi.py
@@ -97,6 +97,8 @@ class MultiLazyBuffer(MathTrait):
   def contiguous(self): return MultiLazyBuffer([x.contiguous() for x in self.lbs], self.axis, self.real)
   def clone(self) -> MultiLazyBuffer: return MultiLazyBuffer([lb.clone() for lb in self.lbs], self.axis, self.real)
   def detach(self) -> MultiLazyBuffer: return MultiLazyBuffer([lb.detach() for lb in self.lbs], self.axis, self.real)
+  @property
+  def toposort(self) -> dict[UOp, None]: return {l:None for x in self.lbs for l in x.toposort}
 
   # elementwise is simple
   def alu(self, op:Ops, *in_srcs:MultiLazyBuffer) -> MultiLazyBuffer:
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 171c7afe31..53ebb9b86c 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -946,6 +946,7 @@ class Tensor(SimpleMathTrait):
       # this is "implicit gradient creation"
       gradient = Tensor(1.0, dtype=self.dtype, device=self.device, requires_grad=False)
 
+    toposort_uop = self.lazydata.toposort
     assert self.shape == gradient.shape, f"grad shape must match tensor shape, {gradient.shape!r} != {self.shape!r}"
     self.grad = gradient
     for t0 in reversed(toposorted):
@@ -958,6 +959,8 @@ class Tensor(SimpleMathTrait):
       for t, g in zip(t0._ctx.parents, grads):
         if g is not None and t.requires_grad:
           assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
+          assert t.lazydata in toposort_uop or (isinstance(t.lazydata, MultiLazyBuffer) and any(x in toposort_uop for x in t.lazydata.lbs)), \
+            f"grad uop must have a path from self\ngrad uop: {t.lazydata}"
           t.grad = g if t.grad is None else (t.grad + g)
       if not retain_graph: del t0._ctx
     return self