diff --git a/test/unit/test_gradient.py b/test/unit/test_gradient.py index e984394ccc..2731253093 100644 --- a/test/unit/test_gradient.py +++ b/test/unit/test_gradient.py @@ -68,6 +68,14 @@ class TestTensorGradient(unittest.TestCase): np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0]) self.assertIs(x.grad, old_grad) + def test_gradient_through_chained_unrealized_setitem(self): + g1 = Tensor.zeros(4).contiguous() + g1[2] = Tensor(1.0) + g2 = Tensor.zeros(5, 4).contiguous() + g2[0] = g1 + x = Tensor.randn(4, 4) + np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4))) + class TestViewGradient(unittest.TestCase): def test_expand(self): x = Tensor.randn(5,2) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 601f9dfab8..2ddde4ecff 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -1,6 +1,6 @@ from dataclasses import dataclass, field, replace import itertools -from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace +from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace, Invalid from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, should_resolve_call from tinygrad.uop.symbolic import symbolic @@ -229,8 +229,9 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp): # if it makes it here, the bufferize is removed # this is the ranges replaced - # NOTE: if buf src is a const, we don't replace it - return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST}, extra_pm=pm_gate_substitute) + # NOTE: if buf src is a const, we don't replace it. if idx is Invalid (dead load), don't replace it either + replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.arg is Invalid)} + return src.substitute(replaced, extra_pm=pm_gate_substitute) def remove_noop_bufferize(idx,b2): if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.BUFFER_VIEW: return None