[PYTHON] Fix PEP 563 compat and remove args_converter#18846
[PYTHON] Fix PEP 563 compat and remove args_converter#18846tqchen wants to merge 2 commits intoapache:mainfrom
Conversation
This PR resolves string annotations from PEP 563 (from __future__ import annotations) at decoration time using typing.get_type_hints(), preventing TypeError crashes when @type_checked passes stringified annotations to isinstance().
This PR removes the @args_converter.auto decorator that auto-converted Python types to Relax Expr by inspecting annotations at decoration time. This was PEP 563-incompatible (annotations become strings). The FFI/C++ layer handles type conversion directly, with minimal explicit conversion added at call sites for Python tuples/lists and primitive types.
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 significantly refactors the argument conversion and type checking mechanisms within the Relax framework. It enhances compatibility with modern Python type annotation practices by correctly resolving postponed evaluations and replaces an implicit, decorator-based argument conversion system with explicit function calls. This change aims to improve the robustness, clarity, and maintainability of the codebase by making argument handling more transparent and less reliant on automatic inference. 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
Activity
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 is a large refactoring that removes the @args_converter.auto decorator and replaces it with explicit calls to convert_to_expr, and also fixes PEP 563 compatibility in the @type_checked decorator. The changes are generally well-executed. However, I've identified a couple of potential bugs in the distributed operators where arguments from Python lists/tuples are not converted to relax.Expr, which could cause runtime errors. I've also found numerous instances where the new explicit conversion logic is more complex than necessary and have suggested simplifications to improve code readability and maintainability. The rest of the changes look good.
| """ | ||
| if isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore | ||
| if isinstance(args, tuple | list): | ||
| args = RxTuple(list(args)) |
There was a problem hiding this comment.
When args is a tuple or list, its elements are not being converted to relax.Expr before being passed to RxTuple. This can lead to a TypeError if args contains non-Expr elements like integers or floats. You can use convert_to_expr to handle this conversion, which will also create the relax.Tuple. You will also need to add from ...utils import convert_to_expr at the top of the file.
| args = RxTuple(list(args)) | |
| args = convert_to_expr(args) |
|
|
||
| if isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore | ||
| if isinstance(args, tuple | list): | ||
| args = RxTuple(list(args)) |
There was a problem hiding this comment.
When args is a tuple or list, its elements are not being converted to relax.Expr before being passed to RxTuple. This can lead to a TypeError if args contains non-Expr elements like integers or floats. You can use convert_to_expr to handle this conversion, which will also create the relax.Tuple. You will also need to add from tvm.relax.utils import convert_to_expr at the top of the file.
| args = RxTuple(list(args)) | |
| args = convert_to_expr(args) |
| if not isinstance(shape, Expr): | ||
| shape = convert_to_expr(shape) |
There was a problem hiding this comment.
This check is redundant because convert_to_expr already handles Expr inputs by returning them as is. You can simplify this by calling convert_to_expr unconditionally.
shape = convert_to_expr(shape)| args = tuple( | ||
| convert_to_expr(a) if isinstance(a, int | float | str | tuple | list) else a | ||
| for a in args | ||
| ) |
There was a problem hiding this comment.
Similar to a previous comment, this argument conversion logic can be simplified. convert_to_expr is idempotent for Expr inputs, so you can apply it to all arguments unconditionally.
args = tuple(convert_to_expr(a) for a in args)| axes = convert_to_expr(axes) if not isinstance(axes, Expr) else axes | ||
| begin = convert_to_expr(begin) if not isinstance(begin, Expr) else begin | ||
| end = convert_to_expr(end) if not isinstance(end, Expr) else end | ||
| if strides is not None and not isinstance(strides, Expr): | ||
| strides = convert_to_expr(strides) |
There was a problem hiding this comment.
This conversion logic can be simplified. convert_to_expr is idempotent on Expr inputs, so the isinstance checks are not necessary.
| axes = convert_to_expr(axes) if not isinstance(axes, Expr) else axes | |
| begin = convert_to_expr(begin) if not isinstance(begin, Expr) else begin | |
| end = convert_to_expr(end) if not isinstance(end, Expr) else end | |
| if strides is not None and not isinstance(strides, Expr): | |
| strides = convert_to_expr(strides) | |
| axes = convert_to_expr(axes) | |
| begin = convert_to_expr(begin) | |
| end = convert_to_expr(end) | |
| if strides is not None: | |
| strides = convert_to_expr(strides) |
| if not isinstance(shape, Expr): | ||
| shape = convert_to_expr(shape) |
There was a problem hiding this comment.
This check is redundant because convert_to_expr already handles Expr inputs by returning them as is. You can simplify this by calling convert_to_expr unconditionally.
shape = convert_to_expr(shape)| if not isinstance(func, Expr): | ||
| func = convert_to_expr(func) |
There was a problem hiding this comment.
This check is redundant because convert_to_expr already handles Expr inputs by returning them as is. You can simplify this by calling convert_to_expr unconditionally.
func = convert_to_expr(func)| args = tuple( | ||
| convert_to_expr(a) if isinstance(a, int | float | str | tuple | list) else a | ||
| for a in args | ||
| ) |
There was a problem hiding this comment.
This argument conversion logic can be simplified. convert_to_expr handles Expr inputs by returning them directly, so you don't need to conditionally call it. Applying it to all arguments is simpler and has the same effect.
args = tuple(convert_to_expr(a) for a in args)| args = py_tuple( | ||
| convert_to_expr(a) if isinstance(a, _CONVERTIBLE_TYPES) else a for a in args | ||
| ) |
There was a problem hiding this comment.
This argument conversion logic can be simplified. convert_to_expr handles Expr inputs by returning them directly, and also handles all types in _CONVERTIBLE_TYPES, so you don't need to conditionally call it. Applying it to all arguments is simpler and has the same effect.
args = py_tuple(convert_to_expr(a) for a in args)| args = py_tuple( | ||
| convert_to_expr(a) if isinstance(a, _CONVERTIBLE_TYPES) else a for a in args | ||
| ) |
There was a problem hiding this comment.
Similar to a previous comment, this argument conversion logic can be simplified. convert_to_expr is idempotent for Expr inputs, so you can apply it to all arguments unconditionally.
args = py_tuple(convert_to_expr(a) for a in args)
1 participant
Summary
Fix @type_checked for PEP 563 compatibility and remove @args_converter.auto decorator.
Changes