From 4b18c92bc5c4a196fa6d7f78bbf60bacc3225cdc Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 18 Jan 2026 00:38:59 -0500 Subject: [PATCH] simpler Context.__enter__ [pr] (#14201) --- tinygrad/helpers.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 49081a8a0b..dfcea9761d 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -148,10 +148,8 @@ def stderr_log(msg:str): print(msg, end='', file=sys.stderr, flush=True) class Context(contextlib.ContextDecorator): def __init__(self, **kwargs): self.kwargs = kwargs def __enter__(self): - self.old_context:dict[str, int] = {} - for k,v in self.kwargs.items(): - self.old_context[k] = ContextVar._cache[k].value - ContextVar._cache[k].value = v + self.old_context:dict[str, int] = {k: ContextVar._cache[k].value for k in self.kwargs} + for k,v in self.kwargs.items(): ContextVar._cache[k].value = v def __exit__(self, *args): for k,v in self.old_context.items(): ContextVar._cache[k].value = v