diff --git a/test/test_schedule.py b/test/test_schedule.py
index c20628b2ff..003ab9423f 100644
--- a/test/test_schedule.py
+++ b/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):
diff --git a/tinygrad/gradient.py b/tinygrad/gradient.py
index 2679775a2e..4f92f259a7 100644
--- a/tinygrad/gradient.py
+++ b/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