[TVMScript] Fix PEP 563 closure variable resolution#18856
Conversation
…lback With `from __future__ import annotations`, Python stores annotations as strings and does not capture annotation-only variables in `__closure__`. This broke TVMScript when buffer shapes/dtypes referenced closure variables. Fix: wrap `extra_vars` in a `collections.ChainMap` with snapshots of all live caller-frame locals (from `inspect.stack()`) as fallback layers in both `tir/entry.py` (`prim_func`) and `ir/entry.py` (`ir_module`). The `ir_module` function now also captures `outer_stack = inspect.stack()` at its entry point, mirroring the existing pattern in `prim_func`. Lookup falls back to frame locals only on cache miss, preserving existing behavior for non-PEP-563 code. Add `tests/python/tvmscript/test_tvmscript_pep563_closure.py` (requires `from __future__ import annotations` at the top) covering closure variables in buffer shapes, dtypes, nested scopes, ir_module, and mixed annotation+body use.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves a critical compatibility issue with Python's PEP 563 ( Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly addresses an issue with PEP 563 closure variable resolution in TVMScript. The approach of using inspect.stack and collections.ChainMap to create a fallback for resolving variables from caller frames is sound. The new test suite provides excellent coverage for various scenarios, including nested closures and usage in IR modules. My review includes a few suggestions to improve code style and consistency.
The previous with_caller_frame_fallback added ALL caller-frame locals to the TVMScript namespace, causing name collisions (e.g., `gv` from a test function shadowing a nested @R.function def gv). Replace with resolve_closure_vars: parses the class source AST to find names used in function annotations, then looks up only those names in enclosing frames. This resolves PEP 563 closure variables (like M in T.Buffer((M,), "float32")) without polluting the namespace.
…cals The previous ChainMap approach added ALL caller-frame locals to the TVMScript namespace, causing name collisions (e.g. `gv` from a test function shadowing a nested @R.function def gv). Replace with resolve_closure_vars: parses the source AST to find names used in function annotations, then looks up only those names in enclosing frames. This resolves PEP 563 closure variables (like M in T.Buffer((M,), "float32")) without polluting the namespace. Applied to both @T.prim_func and @I.ir_module paths.
- Only trigger resolve_closure_vars when annotations are actually strings (PEP 563 active), skipping the AST parse for normal code. - Use __qualname__ to identify lexically enclosing scopes instead of walking the entire caller stack, preventing false matches from unrelated caller frames.
2 participants
With
from __future__ import annotations, Python stores annotations as stringsand does not capture annotation-only variables in
__closure__. This brokeTVMScript when buffer shapes/dtypes referenced closure variables.
Fix: wrap
extra_varsin acollections.ChainMapwith snapshots of all livecaller-frame locals (from
inspect.stack()) as fallback layers in bothtir/entry.py(prim_func) andir/entry.py(ir_module). Their_modulefunction now also captures
outer_stack = inspect.stack()at its entry point,mirroring the existing pattern in
prim_func. Lookup falls back to frame localsonly on cache miss, preserving existing behavior for non-PEP-563 code.
Add
tests/python/tvmscript/test_tvmscript_pep563_closure.py(requiresfrom __future__ import annotationsat the top) covering closure variables inbuffer shapes, dtypes, nested scopes, ir_module, and mixed annotation+body use.