fix UPat.location after pickle (tinygrad#9763)

* fix UPat.location after pickle [pr]

* named upat test

Author: Qazalin

test/test_pickle.py

@@ -21,7 +21,9 @@ def test_pickle_pattern_matcher(self):
 
   def test_pickle_main_pattern_matcher(self):
     from tinygrad.codegen.devectorizer import sym
-    pickle.dumps(sym)
+    ssym = pickle.dumps(sym)
+    dsym = pickle.loads(ssym)
+    self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)
 
   def test_pickle_realized_tensor(self):
     print("** init")

test/test_uops.py

@@ -485,6 +485,9 @@ def test_location(self):
     self.assertEqual(spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
     test_upat = UPat(Ops.CONST, dtypes.bool)
     self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])
+    test_upat_named = test_upat.named("test_name")
+    self.assertEqual(test_upat.location[0], test_upat_named.location[0])
+    self.assertNotEqual(test_upat.location[1], test_upat_named.location[1])
 
 class TestUopsObject(unittest.TestCase):
   # LOL, running this test breaks all instances of "4"

tinygrad/ops.py

@@ -708,7 +708,7 @@ class UPat(MathTrait):
   __slots__ = ("op", "dtype", "arg", "name", "src")
   def __init__(self, op:Optional[Union[Ops, tuple[Ops, ...], set[Ops]]]=None, dtype:Optional[Union[DType, tuple[DType, ...]]]=None,
                src:Optional[Union[tuple[UPat, ...], list[UPat], UPat]]=None, arg:Any=None,
-               name:Optional[str]=None, allow_any_len:bool=False, location=None, custom_early_reject:Optional[set[Ops]]=None):
+               name:Optional[str]=None, allow_any_len:bool=False, custom_early_reject:Optional[set[Ops]]=None, location=None):
     assert op is None or isinstance(op, (Ops, tuple, set)), "op must be Ops or tuple of Ops"
     self.op: Optional[tuple[Ops, ...]] = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op)
     self.dtype: Optional[tuple[DType, ...]] = (dtype,) if isinstance(dtype, DType) else dtype
@@ -733,7 +733,8 @@ def __init__(self, op:Optional[Union[Ops, tuple[Ops, ...], set[Ops]]]=None, dtyp
       upat_match = [src] if isinstance(src, UPat) else ([] if src is None else self.src[0])
       self.early_reject = {pp.op[0] for pp in upat_match if pp.op is not None and len(pp.op) == 1}
 
-  def __reduce__(self): return UPat,(self.op, self.dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject)
+  def __reduce__(self):
+    return UPat, (self.op, self.dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject, self.location)
   def named(self, name:str): return UPat(self.op, self.dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject)
 
   @staticmethod