This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 59a7b0246 [FIX] Defer eager generic tensor annotations (#114)
59a7b0246 is described below
commit 59a7b0246b596316d56f8108ad15eb1ecefb3b80
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 28 13:05:28 2026 -0400
[FIX] Defer eager generic tensor annotations (#114)
---
python/tvm_tirx/tirx/tensor.py | 14 +++++++++++++-
tests/python/tirx/base/test_tensor_type.py | 20 ++++++++++++++++++++
tests/python/tirx/script/test_dynamic_shape.py | 6 ++++--
3 files changed, 37 insertions(+), 3 deletions(-)
diff --git a/python/tvm_tirx/tirx/tensor.py b/python/tvm_tirx/tirx/tensor.py
index a7ced4fdf..28db73cf5 100644
--- a/python/tvm_tirx/tirx/tensor.py
+++ b/python/tvm_tirx/tirx/tensor.py
@@ -18,7 +18,7 @@
from collections.abc import Callable
from numbers import Integral
-from typing import ClassVar
+from typing import ClassVar, TypeVar
import tvm_ffi
@@ -96,6 +96,18 @@ class TensorType(Type):
Strides and element offsets are described by ``layout``.
``align`` is an alias for ``data_alignment``.
"""
+ # Python 3.12/3.13 evaluate generic annotations before the decorator.
+ # Leave a null wrapper until the parser re-evaluates the annotation
+ # source with real symbolic IR variables.
+ dimensions = (shape,) if isinstance(shape, TypeVar) else shape
+ if isinstance(dimensions, (tuple, list)) and any(
+ isinstance(dim, TypeVar) for dim in dimensions
+ ):
+ from tvm_tirx.script.ir_builder.base import IRBuilder
+
+ if not IRBuilder.is_in_scope():
+ return
+
from tvm_tirx.ir.type import _normalize_type
from tvm_tirx.script.ir_builder.base import MISSING, at
diff --git a/tests/python/tirx/base/test_tensor_type.py
b/tests/python/tirx/base/test_tensor_type.py
index 75f2562de..2107167dd 100644
--- a/tests/python/tirx/base/test_tensor_type.py
+++ b/tests/python/tirx/base/test_tensor_type.py
@@ -16,6 +16,7 @@
# under the License.
import uuid
+from typing import TypeVar
import pytest
import tvm_ffi
@@ -24,6 +25,25 @@ import tvm_tirx
import tvm_tirx.testing
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import tirx as T
+from tvm_tirx.script.ir_builder import IRBuilder
+
+
+def test_tensor_type_eager_annotation_placeholder():
+ n = TypeVar("n")
+ assert T.Tensor is tvm_tirx.tirx.TensorType
+ for shape in (n, (n,), [n, 4]):
+ annotation = T.Tensor(shape, T.float32)
+ assert isinstance(annotation, tvm_tirx.tirx.TensorType)
+ assert annotation.__chandle__() == 0
+ with IRBuilder(), pytest.raises(TypeError, match=r"Expected
`Array<ir\.PrimExpr>`"):
+ T.Tensor((n,), T.float32)
+ with pytest.raises(TypeError, match=r"Expected `Array<ir\.PrimExpr>`"):
+ T.Tensor((object(),), T.float32)
+ with pytest.raises(TypeError):
+ T.Tensor(None, T.float32)
+ # Native reconstruction allocates the Python wrapper without constructor
arguments.
+ ty = T.Tensor((4,), T.float32)
+ tvm_tirx.ir.assert_structural_equal(ty,
tvm_tirx.ir.load_json(tvm_tirx.ir.save_json(ty)))
def test_tensor_pointer_type_derived_from_dtype_and_scope():
diff --git a/tests/python/tirx/script/test_dynamic_shape.py
b/tests/python/tirx/script/test_dynamic_shape.py
index a1de09728..6702d5f61 100644
--- a/tests/python/tirx/script/test_dynamic_shape.py
+++ b/tests/python/tirx/script/test_dynamic_shape.py
@@ -127,8 +127,10 @@ def
test_tir_return_annotation_does_not_define_symbolic_var():
@pytest.mark.skipif(sys.version_info < (3, 12), reason="PEP 695 requires
Python 3.12")
@pytest.mark.parametrize("type_param", ["n", "n: T.int32"])
-def test_generic_function_decorators(tmp_path, type_param):
- source = f"""from __future__ import annotations
[email protected]("postponed", [False, True])
+def test_generic_function_decorators(tmp_path, type_param, postponed):
+ source = "from __future__ import annotations\n" if postponed else ""
+ source += f"""
@T.function
def standalone[{type_param}](A: T.Tensor((n,), T.float32)):
T.evaluate(n)
