From 750696a0269d87f09f7d95da71b71f9ea7dc3a7e Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Sun, 8 Sep 2024 17:52:51 +0800
Subject: [PATCH] use uop vars for ast DEFINE_VAR [run_process_replay] (#6404)

prereq for UOps.VALID.
---
 tinygrad/ops.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index 17840b5a12..0cef3a66ac 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -388,8 +388,7 @@ class UOp(MathTrait):
   def vars(self) -> Set[UOp]: return set([x for x in self.sparents if x.op is UOps.DEFINE_VAR])
   def variables(self) -> List[Variable]:
     st_vars: List[Set[Variable]] = [x.st_arg.vars() for x in self.sparents if x.op in BUFFER_UOPS]
-    return sorted(set.union(*st_vars, set([Variable(x.arg[0], x.arg[1], x.arg[2]) for x in self.sparents if x.op is UOps.DEFINE_VAR])),
-      key=lambda v: v.expr)
+    return sorted(set.union(*st_vars, [Variable(x.arg[0], x.arg[1], x.arg[2]) for x in self.vars()]), key=lambda v: v.expr)
   def const_factor(self) -> int:
     """largest known int that divides self"""
     if self.op is UOps.CONST: return self.arg