From c5029fa460e1024779ff1eeeaa6f1fd5cd25c57d Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Sat, 21 Feb 2026 15:33:55 +0800
Subject: [PATCH] jit case with Tensor.empty input, realized means allocated
 (#14930)

* simple failing jit test case with Tensor.empty

* this used to exist in ops.py...

* Revert "removed if self.buffer.is_allocated() in realized (#14836)"

This reverts commit 72cf603805f15bfe7c0372d798e0c89da192c1c1.
---
 test/backend/test_jit.py             | 9 +++++++++
 test/null/test_schedule.py           | 4 ++--
 test/unit/test_realize_is_realize.py | 4 ++--
 test/unit/test_setitem_schedule.py   | 3 ++-
 tinygrad/uop/ops.py                  | 3 ++-
 5 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/test/backend/test_jit.py b/test/backend/test_jit.py
index 91d0f67b84..e440956e22 100644
--- a/test/backend/test_jit.py
+++ b/test/backend/test_jit.py
@@ -490,6 +490,15 @@ class TestJit(unittest.TestCase):
     #with self.assertRaises(JitError):
     #  f(Tensor([2.0])).item()
 
+  def test_jit_init_empty_alt(self):
+    @TinyJit
+    def f(a:Tensor, b:Tensor) -> Tensor: return b.assign(a+1)
+    a = Tensor([1])
+    for _ in range(4):
+      b = Tensor.empty_like(a)
+      c = f(a, b)
+      self.assertEqual(c.item(), 2)
+
 @unittest.skip("Pending multioutput implementation #3607")
 class TestMultioutputJit(unittest.TestCase):
   def _test(self, f):
diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py
index d53c64ca68..fb89f97f7a 100644
--- a/test/null/test_schedule.py
+++ b/test/null/test_schedule.py
@@ -169,7 +169,7 @@ class TestSchedule(unittest.TestCase):
   def test_empty_is_not_realized(self):
     a = Tensor.empty(10)
     child = a+2
-    assert a.uop.is_realized
+    assert not a.uop.is_realized
     child.realize()
     assert a.uop.is_realized
 
@@ -185,7 +185,7 @@ class TestSchedule(unittest.TestCase):
   def test_childless_empty_never_allocates(self):
     a = Tensor.empty(10)
     a.realize()
-    assert not a.uop.buffer.is_allocated()
+    assert not a.uop.is_realized
 
   def test_simplify_padded_const(self):
     a, _ = Tensor.empty(1022).cummax(axis=0)
diff --git a/test/unit/test_realize_is_realize.py b/test/unit/test_realize_is_realize.py
index bfa20a395a..5bf6581269 100644
--- a/test/unit/test_realize_is_realize.py
+++ b/test/unit/test_realize_is_realize.py
@@ -30,14 +30,14 @@ class TestRealizeIsRealized(unittest.TestCase):
 
   def test_empty(self):
     t = Tensor.empty(4, 4).realize()
-    assert t.uop.is_realized
+    assert not t.uop.is_realized
 
   def test_disk(self):
     with tempfile.NamedTemporaryFile() as f:
       f.write(b'\x00' * 16)
       f.flush()
       t = Tensor.empty(4, dtype=dtypes.float32, device=f"disk:{f.name}").realize()
-      assert t.uop.is_realized
+      assert not t.uop.is_realized
 
   def test_assign(self):
     t = Tensor([1, 2, 3])
diff --git a/test/unit/test_setitem_schedule.py b/test/unit/test_setitem_schedule.py
index 34d9be4b18..ac63ef1162 100644
--- a/test/unit/test_setitem_schedule.py
+++ b/test/unit/test_setitem_schedule.py
@@ -36,7 +36,8 @@ class TestSetitemInto(unittest.TestCase):
     self.assertEqual(GlobalCounters.kernel_count, 0)
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4)
+    # TODO: this can be just 4 if empty goes through is_realized setitem path
+    self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly
     t[1].realize()
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)
diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 387c12c2a9..00d31e3955 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -694,7 +694,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op not in (Ops.BUFFER, Ops.MSTACK): return None
     # LUNIQUEs are never realized
     if self.op_in_backward_slice_with_self(Ops.LUNIQUE): return None
-    return self.buffer
+    # NOTE: this is used by the JIT to determine which inputs we capture
+    return self.buffer if self.buffer.is_allocated() else None
   @property
   def is_realized(self) -> bool: return self.base.realized is not None