fix kernelize usage with pm_gradient (#9953)

* fix kernelize usage with pm_gradient

* remove that

test/test_schedule.py

@@ -612,16 +612,25 @@ class TestSchedule(unittest.TestCase):
     e = c.kernelize()+d.kernelize()
     check_schedule(e, 3)
 
-  @unittest.expectedFailure # TODO: this should pass
   def test_kernelize_bw(self):
     a = Tensor.full((3,), 2.0, requires_grad=True).contiguous()
     b = Tensor.full((3,), 3.0, requires_grad=True).contiguous()
     x = (a*b).kernelize()
     y = Tensor.eye(3, requires_grad=True)
     z = y.matmul(x).sum()
-    if getenv("VIZ"):
-      graph_rewrite(z.lazydata, PatternMatcher([]), name="y.matmul(x).sum()")
     z.backward()
+    self.assertEqual(z.item(), 18.0)
+    self.assertEqual(z.grad.item(), 1.0)
+
+  def test_kernelize_bw_view(self):
+    a = Tensor.full((3,1), 2.0, requires_grad=True).contiguous()
+    b = Tensor.full((3,1), 3.0, requires_grad=True).contiguous()
+    x = (a*b).kernelize()
+    y = Tensor.eye(6, requires_grad=True)
+    z = y.matmul(x.expand(3,2).reshape(6)).sum()
+    z.backward()
+    self.assertEqual(z.item(), 36.0)
+    self.assertEqual(z.grad.item(), 1.0)
 
   @unittest.skip("no longer supported")
   def test_double_from(self):

tinygrad/gradient.py

@@ -51,7 +51,7 @@ def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:
   def is_in_target_path(x:UOp) -> bool: return any(u in targets or is_in_target_path(u) for u in x.src) # noqa: F821
   def _walk(node:UOp, visited:set[UOp]) -> Iterator[UOp]:
     visited.add(node)
-    if node.op is Ops.DETACH: return
+    if node.op in {Ops.DETACH, Ops.ASSIGN}: return
     if is_in_target_path(node): # noqa: F821
       for i in node.src:
         if i not in visited: yield from _walk(i, visited) # noqa: F821