This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tvmscript-ast-only-transpiler
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/tvmscript-ast-only-transpiler
by this push:
new 92aeec4d07 [TVMScript] Capture opaque lexical callables through
builders
92aeec4d07 is described below
commit 92aeec4d075b9ab2482f7e2bd53e31f7184230d8
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Sep 21 23:28:31 2026 +0000
[TVMScript] Capture opaque lexical callables through builders
Preserve opaque lexical spellings as metadata placeholders so function
bodies capture callable references without exposing IR values to the
transpiler. Recursive references retain their finalized function binding.
Construction-time source provenance for void and inline calls remains
unresolved.
---
python/tvm/script/ir_builder/construction.py | 2 +-
python/tvm/script/parser_v2/frontend.py | 15 +++++++++++----
python/tvm/script/parser_v2/transpile.py | 17 +++++++++++++----
3 files changed, 25 insertions(+), 9 deletions(-)
diff --git a/python/tvm/script/ir_builder/construction.py
b/python/tvm/script/ir_builder/construction.py
index 6e3e76265c..e362580290 100644
--- a/python/tvm/script/ir_builder/construction.py
+++ b/python/tvm/script/ir_builder/construction.py
@@ -374,7 +374,7 @@ class FunctionRecord:
return self.symbols.symbols[name]
def capture(self, name, fallback):
- """Resolve a free annotation name against retained signature symbols.
+ """Resolve a free lexical name against retained signature symbols.
Parameters
----------
diff --git a/python/tvm/script/parser_v2/frontend.py
b/python/tvm/script/parser_v2/frontend.py
index 1243dc2a36..a523c8dfcd 100644
--- a/python/tvm/script/parser_v2/frontend.py
+++ b/python/tvm/script/parser_v2/frontend.py
@@ -456,13 +456,20 @@ class Compiler:
Only host namespace and callable metadata are exposed to the
transformer.
The name allocator is shared with this compiler; no IR state is shared.
"""
+ # Retain lexical spellings even when their values are opaque. The
+ # transformer may arrange builder-owned captures for those names, but
+ # it receives no concrete value or IR identity for them. A None entry
+ # also prevents a host binding from being mistaken for a builtin.
metadata = {
name: value
+ if (
+ inspect.ismodule(value)
+ or inspect.isfunction(value)
+ or inspect.isclass(value)
+ or type(value).__module__ == "types"
+ )
+ else None
for name, value in self.env.items()
- if inspect.ismodule(value)
- or inspect.isfunction(value)
- or inspect.isclass(value)
- or type(value).__module__ == "types"
}
return IRBuilderTranspiler(
self.filename,
diff --git a/python/tvm/script/parser_v2/transpile.py
b/python/tvm/script/parser_v2/transpile.py
index 3911a6a605..ba5a73521e 100644
--- a/python/tvm/script/parser_v2/transpile.py
+++ b/python/tvm/script/parser_v2/transpile.py
@@ -127,6 +127,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
Original source filename, used in Python compilation and diagnostics.
environment : dict of str to object
Host namespace and callable bindings for static metadata lookup.
+ Other visible lexical spellings map to None, without retaining values.
No expression evaluation or concrete symbol construction occurs here.
builder_name : str
Collision-free injected alias for the current construction namespace.
@@ -156,7 +157,8 @@ class IRBuilderTranspiler(ast.NodeTransformer):
-----
``filename``, the span and name callbacks, and ``infrastructure_name``
belong
to one source unit. ``namespace_bindings`` maps visible source spellings to
- host namespaces and callables; assignments and parameters shadow entries.
+ host namespaces and callables, or None for opaque lexical values;
+ assignments and parameters shadow entries.
``dialect_prefix`` selects the builder alias and is restored after
functions.
``bound`` holds identifier strings only. ``optional`` maps conditionally
exported names to AST reads of cached named outputs, so a name defined by
@@ -1342,9 +1344,16 @@ class IRBuilderTranspiler(ast.NodeTransformer):
for item in ast.walk(statement)
if isinstance(item, ast.Name) and isinstance(item.ctx, ast.Load)
}
- free_names = annotation_names & body_names -
self._assigned_names(node.body) - {
- arg.arg for arg in node.args.args
- }
+ # Known lexical values used only in the body also cross the existing
+ # builder capture boundary. This lets callable references keep dialect
+ # behavior without inspecting a call target or changing ordinary f(x).
+ # Unknown names retain normal Python lookup/error behavior.
+ free_names = (
+ annotation_names | self.namespace_bindings.keys()
+ ) & body_names - self._assigned_names(node.body) - {arg.arg for arg in
node.args.args}
+ # Recursive references use the factory's finalized function binding,
+ # never an older outer function with the same spelling.
+ free_names.discard(node.name)
for name in sorted(free_names):
arguments.kwonlyargs.append(ast.arg(name))
arguments.kw_defaults.append(
