improve typing for tir.PrimExpr

[clouds56]

Motivation

Now we use PrimExpr as type hint in most places and it would report errors (in static typing) sometimes.

case 1:

a = Var('a', dtype='int')
a + a # this is Ok since __add__(self, other: PrimExpr) -> PrimExpr
a + 1 # this is not Ok
a + a + 1 # this is not Ok since (a+a) is a PrimExpr, not PrimExprWithOp

case 2:

a = Var('a', dtype='int')
T.grid(a, a, a) # this is ok since a is PrimExpr
T.grid(1, 2, 3) # this is not Ok, since 1 is not a PrimExpr

Solution

This PR is to mitigate these cases by introducing

PrimIntExpr = Union[PrimExpr, int]
PrimFloatExpr = Union[PrimExpr, float]
PrimLogicalExpr = Union[PrimeExpr, int, bool]

Problems

• If there's no problem with current idea, I'd like to apply PrimIntExpr/PrimFloatExpr to more places to improve typing system.
• It might be a huge changeset, should I split it into several PR?
• How to handle pylint E0633, upgrade pylint or add some ignore for this?
• ForFrame doesn't have __iter__, this would be fixed in another PR. Current PR would focus on PrimExpr. POC here 7d64de1. It could be included in this PR if you like.

[gemini-code-assist]

Summary of Changes

Hello @clouds56, I'm Gemini Code Assist¹! 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 improves the type annotations within the tir.PrimExpr module. By introducing more precise type hints for binary and unary operations, it allows PrimExpr to interact more seamlessly with standard Python numeric types and provides more specific return types for these operations. This enhances the overall type safety, readability, and maintainability of the PrimExpr API, facilitating better static analysis.

Highlights

• Enhanced Type Hinting: The typing.Self type was imported and utilized to provide more precise return type annotations for methods that return an instance of their own class.
• New Numeric Type Alias: A new type alias numeric = Union[int, float, complex] was introduced to simplify and standardize type hints for operations involving standard Python numeric types.
• Improved Operator Overload Signatures: The magic methods within the ExprOp class (e.g., __add__, __mul__, __eq__) now accept Union[PrimExpr, numeric] (or similar) for the other operand, allowing for more flexible and type-safe interactions with both PrimExpr instances and native Python numbers. Additionally, their return types are now more specific, indicating the exact PrimExpr subclass (e.g., 'Add', 'Mul') that results from the operation.
• Refined astype Return Type: The astype method's return type was updated to Union['Cast', 'Self'], providing a more accurate representation of its potential return values.

🧠 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.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request refines type hints in the ExprOp class within python/tvm/tir/expr.py by broadening the other argument for various operator overload methods (e.g., __add__, __mul__, __div__, comparison, and bitwise operators) to accept numeric types (int, float, complex) or int/bool in addition to PrimExpr. It also makes the return types of these methods more specific, using string literals like "Add", "Mul", "Call", and updates the astype method's return type to Union["Cast", "Self"]. A new numeric type alias was introduced, and Self was imported. A review comment suggests further updating the type hints for __xor__ and __rxor__ to include bool as an accepted type for the other argument, ensuring consistency with __and__ and __or__ given Python's behavior for bitwise operations on booleans.

[clouds56]

The lint is false positive in this case E0633: Attempting to unpack a non-sequence (unpacking-non-sequence), see pylint-dev/pylint#4696

I update pylint from 3.0.0 to 3.3.9 the problem disappears.