diff --git a/test/test_rangeify.py b/test/test_rangeify.py index 8558c333ea..221a87369a 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -192,5 +192,26 @@ class TestOuterworld(unittest.TestCase): self.assertTrue((manual==out).all().item()) + def test_setitem_pyrange(self): + with Context(DEBUG=0): + t = Tensor.rand(10).realize() + o = Tensor.empty(10) + GlobalCounters.reset() + for i in range(10): + o[i] = t[i] + o.realize() + self.assertTrue((t==o).all().item()) + + @unittest.skip("TODO: fix this") + def test_setitem(self): + with Context(DEBUG=0): + t = Tensor.rand(10).realize() + o = Tensor.empty(10) + GlobalCounters.reset() + i = UOp.range(10, -1) + o[i] = t[i] + o.contiguous(i).realize() + self.assertTrue((t==o).all().item()) + if __name__ == '__main__': unittest.main()