diff --git a/test/backend/test_pickle.py b/test/backend/test_pickle.py index cbb0aa3782..7033a4dd6b 100644 --- a/test/backend/test_pickle.py +++ b/test/backend/test_pickle.py @@ -2,7 +2,7 @@ import unittest, pickle, types, tracemalloc import numpy as np from tinygrad import Tensor, Device, TinyJit, Variable, dtypes from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV -from tinygrad.uop.ops import PatternMatcher, UPat, UOp +from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function class TestPickle(unittest.TestCase): def test_pickle_code_object(self): @@ -11,6 +11,11 @@ class TestPickle(unittest.TestCase): fxn = types.FunctionType(pickle.loads(code_str), globals()) self.assertEqual(fxn(2), 4) + def test_deconstruct_function_nested_comprehension(self): + # pre PEP 709, each comprehension is its own code object, so dtypes here is referenced two code objects deep + def fxn(): return [[dtypes.int for _ in range(2)] for _ in range(2)] + self.assertEqual(types.FunctionType(*deconstruct_function(fxn))(), fxn()) + def test_pickle_pattern_matcher(self): pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)]) sink = UOp.const(2) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index fadcdb866e..6611ef12bf 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1423,9 +1423,9 @@ class UPat(OpMixin): return res def deconstruct_function(fxn:Callable) -> tuple: - new_globals = {k:v for k,v in fxn.__globals__.items() if k in fxn.__code__.co_names} - for co in fxn.__code__.co_consts: - if isinstance(co, types.CodeType): new_globals.update({k:v for k,v in fxn.__globals__.items() if k in co.co_names}) + # globals can be referenced from arbitrarily nested code objects (comprehensions/lambdas, pre PEP 709) + def names(co:types.CodeType) -> set: return set(co.co_names).union(*(names(c) for c in co.co_consts if isinstance(c, types.CodeType))) + new_globals = {k:v for k,v in fxn.__globals__.items() if k in names(fxn.__code__)} # NOTE: optional round trip through pickle! assert fxn.__closure__ is None, "closures are not supported in pattern matchers" ret = fxn.__code__, new_globals, fxn.__name__, fxn.__defaults__