From 55f5eee6b9fc7227f3652138a2523c20350e3178 Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Sun, 22 Feb 2026 11:30:24 -0500 Subject: [PATCH 1/2] one run_schedule in realize --- test/unit/test_assign.py | 5 ++--- tinygrad/tensor.py | 13 ++++++++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py index b39b52c23fea5..7d913d16c926d 100644 --- a/test/unit/test_assign.py +++ b/test/unit/test_assign.py @@ -756,13 +756,12 @@ def test_interleaved_assign_read_patterns(self): np.testing.assert_equal(b.numpy(), [1, 2, 3, 4]) def test_variable_slice_ordering(self): - """Variable-indexed slices - tests symbolic dependency tracking.""" + """Variable-indexed slices with conflicting bindings are not allowed.""" v_i = Variable("i", 0, 3) buf = Tensor.zeros(4, 4).contiguous().realize() buf[v_i.bind(0):v_i.bind(0)+1, :].assign(Tensor.ones(1, 4)) buf[v_i.bind(1):v_i.bind(1)+1, :].assign(Tensor.ones(1, 4) * 2) - self.assertEqual(buf[0:1, :].sum().item(), 4) - self.assertEqual(buf[1:2, :].sum().item(), 8) + with self.assertRaises(RuntimeError): buf.sum().realize() def test_multi_step_assign_read_write_same_buffer(self): """Assign to m and param reading b, then update b, across multiple steps. diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index e46e59d8be761..b183fc05d2381 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -271,16 +271,20 @@ def schedule(self, *lst:Tensor) -> list[ExecItem]: @disable_gc() def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor: """Triggers the computation needed to create these Tensor(s).""" - # side-realize pending assigns for buffers referenced by these tensors + # collect and schedule pending assigns, then schedule main computation, run everything once + all_schedules: list[ExecItem] = [] + all_var_vals: dict[str, int] = {} if _pending_assigns: def _realize_pending(buf): + nonlocal all_var_vals for assign_uop in _pending_assigns.pop(buf, []): # recursively realize pending assigns that this assign's value depends on for u in assign_uop.toposort(): if u.op is Ops.BUFFER and u in _pending_assigns: _realize_pending(u) becomes_map, schedule, var_vals = complete_create_schedule_with_vars(UOp.sink(assign_uop)) _apply_map_to_tensors(becomes_map, name="Apply Pending Assign") - run_schedule(schedule, var_vals, do_update_stats=do_update_stats) + all_schedules.extend(schedule) + all_var_vals = merge_dicts([all_var_vals, var_vals]) # update remaining pending assigns so they reference realized buffers instead of stale lazy graphs if becomes_map: for assigns in _pending_assigns.values(): @@ -288,7 +292,10 @@ def _realize_pending(buf): for buf in {u for t in (self,)+lst for u in t.uop.toposort() if u.op is Ops.BUFFER}: if buf in _pending_assigns: _realize_pending(buf) if len(to_realize:=[x for x in (self,)+lst if not x.uop.has_buffer_identity()]): - run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats) + schedule, var_vals = Tensor.schedule_with_vars(*to_realize) + all_schedules.extend(schedule) + all_var_vals = merge_dicts([all_var_vals, var_vals]) + if all_schedules: run_schedule(all_schedules, all_var_vals, do_update_stats=do_update_stats) return self def replace(self, x:Tensor) -> Tensor: From 927acf73439b231ea3649f3aacce47cd939862c9 Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Sun, 22 Feb 2026 13:00:29 -0500 Subject: [PATCH 2/2] gc? --- tinygrad/tensor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index b183fc05d2381..095bebc1dd22a 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -268,7 +268,6 @@ def schedule(self, *lst:Tensor) -> list[ExecItem]: assert len(var_vals) == 0 return schedule - @disable_gc() def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor: """Triggers the computation needed to create these Tensor(s).""" # collect and schedule pending assigns, then schedule main computation, run everything once