[TIR] 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.

[tlopex]

@clouds56 Thanks for working on this, and sorry for the late reply.

After looking through the full change, I don’t think we should have it in the current form. I’d rather say that clearly now than keep dragging the review out.

Main concerns:

1. Scope does not match the project’s typing discipline.

TVM CI only enforces ruff today; we do not enforce mypy or pyright anywhere in the project. That means adding several new public aliases plus a large number of signature changes increases our maintenance burden, without any mechanism that ensures those annotations remain correct over time. In practice, we would be taking on long-term annotation debt for limited in-tree benefit.

2. It conflicts with TIR’s intended semantic boundary.

In TIR, PrimExpr is the canonical API, while Python int / float are convenience values converted implicitly through the FFI boundary. Annotating parameters everywhere as Union[PrimExpr, int] effectively promotes that convenience path into the type-level contract. That weakens the signal at call sites and makes dtype behavior less explicit. For example, in an expression like a(int32) + 1e12, the annotation no longer encourages the reader to think about what dtype is actually intended. I think we should keep PrimExpr as the documented contract.

3. Some of the narrower return types are not semantically correct.

For example, annotating __add__ as returning "Add" is not generally accurate, since constant folding may produce an IntImm, and mixed-dtype cases may introduce a Cast. Similarly, annotating __eq__ as returning "EqualOp" conflicts with the conventional __eq__ -> bool expectation and is likely to create new checker noise downstream at ordinary == call sites. Also, PrimLogicalExpr = Union[PrimExpr, int, bool] does not really carry additional meaning, since bool is already a subclass of int in Python.

4. This feels like a symptom-level fix rather than a root-level fix.

The motivating issue is real: IDE squiggles on common patterns like a + 1 or T.grid(1, 2, 3) are annoying. But addressing that by widening annotations across a large number of APIs means every future builder or helper will need the same treatment. If we want a durable solution, it should probably be handled closer to the FFI / PrimExpr conversion boundary, rather than being distributed across many individual call sites.

Thanks again for the effort here. I don’t think the overall motivation is wrong, but the current scope and the interaction with TIR’s dtype/semantic discipline make this hard to merge as-is.

[clouds56]

Thanks for the detailed feedback.
I’d like to clarify my thinking and discuss the direction a bit further.

I’m still relatively new to TVM and am actively learning the primitives and APIs, so please feel free to correct me if I misunderstand anything.

annotation in the repo should be more like a hint but not a contract

I understand your concern about maintenance cost. Correct me if I'm wrong, now the typing status of our project is not enforced, the existing typing might behave more like a hint but not a contract.

Everyone uses type hint in different ways, most of them are used for mypy/IDE to check.
Since we don’t enforce mypy, it’s unclear who the primary consumers of these type hints are—users, developers, or both.
My view is that we should optimize for usability here: even if annotations are not perfectly “correct,” making them more kindly still provides practical value.
For example, we annotated __add__ to PrimExpr is correct, while annotated __add__ to PrimExprWithOp/Add is kindly.

we would be taking on long-term annotation debt for limited in-tree benefit.

Now annotations in this repo is already unsound (not matching what it actually is), so maintain of annotation would be definitely a form of debt. I have 3 ways towards future:

1. Gradually move toward more kindly annotations, that would improve IDE support, like auto complete and type check, that would provide in-tree benefits, right?
2. Do minimal maintain for annotations, fix cases like def __eq__(self, other: PrimExpr) -> PrimExpr discussed below to def __eq__(self, other: PrimExpr) -> PrimExpr | bool.
3. Do nothing, leave outdated annotation as is, they are of no value after they are created. This approach increases technical debt.

TIR’s intended semantic boundary.

Similarly, annotating eq as returning "EqualOp" conflicts with the conventional eq -> bool expectation

Taking __eq__ as an example, annotating it as EqualOp is wrong, but the original code def __eq__(self, other: PrimExpr) -> PrimExpr conflicts with bool result as well. I think this wouldn't be likely fixed in short future.
So if both is not correct, EqualOp here might be more intuitive for readers, and more helpful for IDE, since EqualOp is a PrimExprWithOp.
I'd like to know your final decision what signature should be here to update the PR.

And for Union[PrimExpr, int] or PrimIntExpr, this just occurs in final API for input paramters.

This feels like a symptom-level fix rather than a root-level fix.

Yes, your are right. I’m not familiar with tvm's codebase, as an end user, I interact mostly things like @I.ir_module and @R.function that's why I focused on API like buffer, __add__, etc. first.
And in my view, annotate around FFI seems like a more long-term solution, while it might have more work to be done than annotate around user facing API, which is relatively more stable. Internal FFI might be refactored during the time.

I agree there's lots of API needs to be annotated, but we could separate it into multiple steps. Each step is no harm and provide a better future.

roadmap of tvm

I like the idea of tvm, and have been following the repo quite a while, back to the period when it used to be designed to replace the underlying framework of mxnet. I once thought that's the future of foundation of ML infra.
Today I still think tvm is an impressive piece of engineering. However, from a user perspective, it can be difficult to work with today. The lack of up-to-date documentation, limited examples, and weak IDE support (e.g., autocomplete) can discourage adoption.
Frameworks like PyTorch succeed in part because they are highly user-friendly and accessible to a broad audience. More users lead to more potential contributors.

final note

I fully acknowledge that this PR may not be the right shape yet, and it likely contains mistakes.

However, I do believe that improving user-facing type annotations is a worthwhile direction for TVM’s development, and I’d be happy to adjust the scope or approach based on your guidance.

For the prosperity of tvm.

[tlopex]

Thanks for the pushback, @clouds56! I want to revise part of my earlier reply.

I agree that the input-side widening (PrimExpr -> Union[PrimExpr, int, float]) is reasonable. PrimExpr in C++ already has implicit constructors from integer and floating-point literals, and the Python side of TVM also relies heavily on that behavior in practice. So in that sense, the current annotation is narrower than the actual contract. I’m fine with landing that part, along with the PrimIntExpr / PrimFloatExpr aliases.

What I still do not agree with:

• Return narrowing such as -> "Add" is not sound. _generic.add can fold to an IntImm, or introduce a Cast during promotion, so annotating it as always returning Add is more misleading than leaving it as PrimExpr.
• __eq__ -> "EqualOp" should be discussed separately. If we want to annotate it, the choices are more subtle (bool, EqualOp, a union, or leaving it unannotated). I do not think “both are imperfect, so pick the more specific one” is a good justification.
• PrimLogicalExpr = Union[PrimExpr, int, bool] does not look useful as written, since bool is already a subclass of int.

More broadly, I do not think this PR should be treated as approval for the larger direction. Extending the same pattern to more APIs, upgrading pylint expectations, or including unrelated typing cleanups like ForFrame.__iter__ should each be separate decisions and separate PRs.

Whether it still makes sense to push this kind of change — even scoped down to input widening + expression aliases — probably deserves more discussion before we commit to a direction.

cc @tqchen — would appreciate your opinion. Specifically: given that we do not currently enforce mypy / pyright in CI, is it worth introducing PrimIntExpr / PrimFloatExpr and spreading them across TIR APIs? Or would you rather we leave the annotations as-is?