fix Variable init from the DEFINE_VAR refactor (#6448)

prereq for UOps.VALID.

test/test_uops.py

@@ -1,6 +1,7 @@
 from typing import Optional, Tuple, Any, List
 import unittest, math
 import numpy as np
+from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.helpers import CI, DEBUG, getenv, Context
 from tinygrad.dtype import dtypes, DType, PtrDType
@@ -348,8 +349,8 @@ class TestAssembly(unittest.TestCase):
     self.assertEqual(uops[-1].arg, BinaryOps.SHR)
     self.assertEqual(uops[-2].arg, BinaryOps.IDIV)
 
-class TestUOpCompare(unittest.TestCase):
-  def test_alu_same_src_different_arg(self):
+class TestUOpMethod(unittest.TestCase):
+  def test_compare_alu_same_src_different_arg(self):
     a = UOp(UOps.CONST, dtypes.float, (), 2.0)
     b = UOp(UOps.CONST, dtypes.float, (), 3.0)
 
@@ -357,6 +358,15 @@ class TestUOpCompare(unittest.TestCase):
     mul = UOp(UOps.ALU, dtypes.float, (a, b), BinaryOps.MUL)
     assert (add < mul) or (mul < add), "add and mul with same src should have an order"
 
+  def test_uop_variables(self):
+    a = Variable("a", 1, 10)
+    uop_var = UOp.const(dtypes.int, a)
+    st_var = UOp(UOps.LOAD, dtypes.float, (UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), (), 0),
+                                           ShapeTracker.from_shape((2, a)).to_uop()))
+    ast_vars = (st_var+uop_var).variables()
+    self.assertEqual(len(ast_vars), 1)
+    self.assertEqual(ast_vars[0], a)
+
 class TestUOpStr(unittest.TestCase):
   def test_uop_str(self):
     a = UOp(UOps.CONST, dtypes.float, (), 2.0) + UOp(UOps.CONST, dtypes.float, (), 3.0)
@@ -376,6 +386,7 @@ class TestUOpStr(unittest.TestCase):
     assert str(eval(str(a))) == str(a)
 
   def test_variable_const(self):
+    # TODO: this is not possible after VALID.
     uop = UOp(UOps.CONST, dtypes.int, (), arg=Variable("a",1,10))
     assert str(eval(str(uop))) == str(uop)
 

tinygrad/ops.py

@@ -403,7 +403,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, [Variable(x.arg[0], x.arg[1], x.arg[2]) for x in self.vars()]), key=lambda v: v.expr)
+    return sorted(set.union(*st_vars, [Variable(x.arg[0], x.arg[1].arg, x.arg[2].arg) 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