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-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new a6be165d [FEAT] Add fallback types for named schemas (#819)
a6be165d is described below
commit a6be165df636b3e3f173296c3f5f0889042fb82c
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 26 16:58:26 2026 -0400
[FEAT] Add fallback types for named schemas (#819)
## Summary
Preserve a named schema's precise type and optional fallback in Python
TypeSchema metadata. Binding annotations use the fallback when the
primary schema has no supported annotation origin, including inside
nested type constructors.
Runtime type checking and conversion continue to use the primary schema.
---
python/tvm_ffi/core.pyi | 2 ++
python/tvm_ffi/cython/type_info.pxi | 19 +++++++++++++++----
python/tvm_ffi/stub/utils.py | 7 ++++++-
tests/python/test_type_converter.py | 9 +++++++++
4 files changed, 32 insertions(+), 5 deletions(-)
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index 1ebed00d..de5ddac4 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -271,12 +271,14 @@ class TypeSchema:
origin: str
args: tuple[TypeSchema, ...] = ()
origin_type_index: int
+ fallback: TypeSchema | None
def __init__(
self,
origin: str,
args: tuple[TypeSchema, ...] = (),
origin_type_index: int = ...,
+ fallback: TypeSchema | None = None,
) -> None: ...
@staticmethod
def from_json_obj(obj: dict[str, Any]) -> TypeSchema: ...
diff --git a/python/tvm_ffi/cython/type_info.pxi
b/python/tvm_ffi/cython/type_info.pxi
index 02c7c8b1..cc6e0006 100644
--- a/python/tvm_ffi/cython/type_info.pxi
+++ b/python/tvm_ffi/cython/type_info.pxi
@@ -168,6 +168,7 @@ class TypeSchema:
origin: str
args: tuple["TypeSchema", ...] | None = None
origin_type_index: int =
dataclasses.field(default=_ORIGIN_TYPE_INDEX_UNKNOWN, repr=False)
+ fallback: "TypeSchema | None" = None
def __post_init__(self):
origin = self.origin
@@ -270,8 +271,9 @@ class TypeSchema:
)
origin = obj["type"]
origin = _TYPE_SCHEMA_ORIGIN_CONVERTER.get(origin, origin)
+ fallback = TypeSchema.from_json_obj(obj["fallback"]) if "fallback" in
obj else None
if "args" not in obj:
- return TypeSchema(origin)
+ return TypeSchema(origin, fallback=fallback)
raw_args = obj["args"]
if not isinstance(raw_args, (list, tuple)):
raw_args = ()
@@ -279,7 +281,7 @@ class TypeSchema:
TypeSchema.from_json_obj(a) for a in raw_args
if isinstance(a, dict)
)
- return TypeSchema(origin, args)
+ return TypeSchema(origin, args, fallback=fallback)
@staticmethod
def from_json_str(s: str) -> "TypeSchema":
@@ -573,6 +575,11 @@ class TypeSchema:
input_mode: bool,
expanded_convert_types: "frozenset[int]",
) -> str:
+ # A fallback changes only the binding annotation. Keep the primary
+ # schema intact for runtime type checking and conversion.
+ if (self.fallback is not None and self.origin_type_index ==
_ORIGIN_TYPE_INDEX_UNKNOWN
+ and self.origin not in ("tuple", "list", "dict")):
+ return self.fallback._repr_impl(ty_map, input_mode,
expanded_convert_types)
if input_mode and self.origin_type_index >= kTVMFFIStaticObjectBegin:
if self.origin_type_index not in expanded_convert_types:
raw_convert_schema = _lookup_type_attr(
@@ -613,11 +620,15 @@ class TypeSchema:
def to_json(self) -> dict[str, Any]:
"""Convert a TypeSchema to a JSON-compatible dict."""
if self.args is not None and (self.args or self.origin == "tuple"):
- return {
+ result = {
"type": self.origin,
"args": [a.to_json() for a in self.args],
}
- return {"type": self.origin}
+ else:
+ result = {"type": self.origin}
+ if self.fallback is not None:
+ result["fallback"] = self.fallback.to_json()
+ return result
def _annotation_union(args):
diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py
index d6b8bca9..f317f9c4 100644
--- a/python/tvm_ffi/stub/utils.py
+++ b/python/tvm_ffi/stub/utils.py
@@ -131,7 +131,12 @@ class NamedTypeSchema(TypeSchema):
frozen: bool = False,
) -> None:
"""Initialize a `NamedTypeSchema` with the given name, schema and
field facts."""
- super().__init__(origin=schema.origin, args=schema.args)
+ super().__init__(
+ origin=schema.origin,
+ args=schema.args,
+ origin_type_index=schema.origin_type_index,
+ fallback=schema.fallback,
+ )
self.name = name
self.size = size
self.alignment = alignment
diff --git a/tests/python/test_type_converter.py
b/tests/python/test_type_converter.py
index ab582990..16c6c7ab 100644
--- a/tests/python/test_type_converter.py
+++ b/tests/python/test_type_converter.py
@@ -3527,6 +3527,15 @@ class TestTypeSchemaAnnotationRendering:
assert schema.input_repr(_input_ty_map) == "Mapping[str,
Sequence[int]]"
assert schema.repr(_output_ty_map) ==
schema.output_repr(_output_ty_map)
+ def test_named_schema_fallback(self) -> None:
+ schema = TypeSchema.from_json_obj(
+ {
+ "type": "example.PositiveInt",
+ "fallback": {"type": "int"},
+ }
+ )
+ assert schema.output_repr() == "int"
+
def test_object_convert_type_schema_attr_widens_input_only(self) -> None:
"""__ffi_convert_type_schema__ affects input annotations only."""
type_key = _unique_type_key("ConvertTypeSchema")
