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.
Tianqi Chen
61f80814e6