gemini-code-assist[bot] commented on code in PR #18640:
URL: https://github.com/apache/tvm/pull/18640#discussion_r2663921667
##########
python/tvm/tir/expr.py:
##########
@@ -74,111 +77,111 @@ class ExprOp:
# TODO(tkonolige): use inspect to add source information to these objects
- def __add__(self, other: PrimExpr) -> PrimExpr:
+ def __add__(self, other: Union[PrimExpr, numeric]) -> "Add":
return _generic.add(self, other)
- def __radd__(self, other: PrimExpr) -> PrimExpr:
+ def __radd__(self, other: Union[PrimExpr, numeric]) -> "Add":
return _generic.add(other, self)
- def __sub__(self, other: PrimExpr) -> PrimExpr:
+ def __sub__(self, other: Union[PrimExpr, numeric]) -> "Sub":
return _generic.subtract(self, other)
- def __rsub__(self, other: PrimExpr) -> PrimExpr:
+ def __rsub__(self, other: Union[PrimExpr, numeric]) -> "Sub":
return _generic.subtract(other, self)
- def __mul__(self, other: PrimExpr) -> PrimExpr:
+ def __mul__(self, other: Union[PrimExpr, numeric]) -> "Mul":
return _generic.multiply(self, other)
- def __rmul__(self, other: PrimExpr) -> PrimExpr:
+ def __rmul__(self, other: Union[PrimExpr, numeric]) -> "Mul":
return _generic.multiply(other, self)
- def __div__(self, other: PrimExpr) -> PrimExpr:
+ def __div__(self, other: Union[PrimExpr, numeric]) -> "Div":
if _dtype_is_int(self) and _dtype_is_int(other):
raise div_ambiguity_error()
return _generic.divide(self, other)
- def __rdiv__(self, other: PrimExpr) -> PrimExpr:
+ def __rdiv__(self, other: Union[PrimExpr, numeric]) -> "Div":
if _dtype_is_int(self) and _dtype_is_int(other):
raise div_ambiguity_error()
return _generic.divide(other, self)
- def __truediv__(self, other: PrimExpr) -> PrimExpr:
+ def __truediv__(self, other: Union[PrimExpr, numeric]) -> "Div":
if _dtype_is_int(self) and _dtype_is_int(other):
raise div_ambiguity_error()
return _generic.divide(self, other)
- def __rtruediv__(self, other: PrimExpr) -> PrimExpr:
+ def __rtruediv__(self, other: Union[PrimExpr, numeric]) -> "Div":
if _dtype_is_int(self) and _dtype_is_int(other):
raise div_ambiguity_error()
return _generic.divide(other, self)
- def __floordiv__(self, other: PrimExpr) -> PrimExpr:
+ def __floordiv__(self, other: Union[PrimExpr, numeric]) -> "FloorDiv":
return _generic.floordiv(self, other)
- def __rfloordiv__(self, other: PrimExpr) -> PrimExpr:
+ def __rfloordiv__(self, other: Union[PrimExpr, numeric]) -> "FloorDiv":
return _generic.floordiv(other, self, None)
- def __mod__(self, other: PrimExpr) -> PrimExpr:
+ def __mod__(self, other: Union[PrimExpr, numeric]) -> "Mod":
return _ffi_api._OpFloorMod(self, other, None) # type: ignore
- def __rmod__(self, other: PrimExpr) -> PrimExpr:
+ def __rmod__(self, other: Union[PrimExpr, numeric]) -> "Mod":
return _ffi_api._OpFloorMod(other, self, None) # type: ignore
- def __neg__(self) -> PrimExpr:
+ def __neg__(self) -> "Mul":
neg_one = const(-1, self.dtype) # type: ignore
return self.__mul__(neg_one)
- def __lshift__(self, other: PrimExpr) -> PrimExpr:
+ def __lshift__(self, other: Union[PrimExpr, int]) -> "Call":
return _ffi_api.left_shift(self, other, None) # type: ignore
- def __rlshift__(self, other: PrimExpr) -> PrimExpr:
+ def __rlshift__(self, other: Union[PrimExpr, int]) -> "Call":
return _ffi_api.left_shift(other, self, None) # type: ignore
- def __rshift__(self, other: PrimExpr) -> PrimExpr:
+ def __rshift__(self, other: Union[PrimExpr, int]) -> "Call":
return _ffi_api.right_shift(self, other, None) # type: ignore
- def __rrshift__(self, other: PrimExpr) -> PrimExpr:
+ def __rrshift__(self, other: Union[PrimExpr, int]) -> "Call":
return _ffi_api.right_shift(other, self, None) # type: ignore
- def __and__(self, other: PrimExpr) -> PrimExpr:
+ def __and__(self, other: Union[PrimExpr, int, bool]) -> "Call":
return _ffi_api.bitwise_and(self, other, None) # type: ignore
- def __rand__(self, other: PrimExpr) -> PrimExpr:
+ def __rand__(self, other: Union[PrimExpr, int, bool]) -> "Call":
return _ffi_api.bitwise_and(other, self, None) # type: ignore
- def __or__(self, other: PrimExpr) -> PrimExpr:
+ def __or__(self, other: Union[PrimExpr, int, bool]) -> "Call":
return _ffi_api.bitwise_or(self, other, None) # type: ignore
- def __ror__(self, other: PrimExpr) -> PrimExpr:
+ def __ror__(self, other: Union[PrimExpr, int, bool]) -> "Call":
return _ffi_api.bitwise_or(other, self, None) # type: ignore
- def __xor__(self, other: PrimExpr) -> PrimExpr:
+ def __xor__(self, other: Union[PrimExpr, int]) -> "Call":
return _ffi_api.bitwise_xor(self, other, None) # type: ignore
- def __rxor__(self, other: PrimExpr) -> PrimExpr:
+ def __rxor__(self, other: Union[PrimExpr, int]) -> "Call":
return _ffi_api.bitwise_xor(other, self, None) # type: ignore
Review Comment:
For consistency with `__and__` and `__or__`, which accept `bool` as a valid
type for the `other` argument, `__xor__` and `__rxor__` should also be updated
to accept `bool`. In Python, bitwise operations on booleans are well-defined,
where they are treated as integers (0 or 1).
```suggestion
def __xor__(self, other: Union[PrimExpr, int, bool]) -> "Call":
return _ffi_api.bitwise_xor(self, other, None) # type: ignore
def __rxor__(self, other: Union[PrimExpr, int, bool]) -> "Call":
return _ffi_api.bitwise_xor(other, self, None) # type: ignore
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
