minor update

Author: kurisu6912

tilelang/language/tir/ir.pyi

@@ -51,12 +51,12 @@ def nextafter(x1: _T, x2: _T) -> _T: ...
 def popcount(x: _T) -> _T: ...
 def pow(x: _T, y: _T, span: Span | None=None) -> _T: ...
 def q_multiply_shift(x: _T, y: _T, q: _T, s: _T) -> _T: ...
-def q_multiply_shift_per_axis(x: PrimExpr, y: PrimExpr, ls: PrimExpr, rs: PrimExpr, q: IntImm, is_lshift_required: IntImm, is_rshift_required: IntImm) -> PrimExpr: ...
+def q_multiply_shift_per_axis(x: _T, y: _T, ls: _T, rs: _T, q: IntImm, is_lshift_required: IntImm, is_rshift_required: IntImm) -> PrimExpr: ...
 def ret(val: _T) -> _T: ...
 def round(x: _T, span: Span | None=None) -> _T: ...
 def rsqrt(x: _T) -> _T: ...
-def shift_left(x: PrimExpr, y: PrimExpr, span=None) -> PrimExpr: ...
-def shift_right(x: PrimExpr, y: PrimExpr, span=None) -> PrimExpr: ...
+def shift_left(x: _T, y: _T, span=None) -> _T: ...
+def shift_right(x: _T, y: _T, span=None) -> _T: ...
 def sigmoid(x: _T) -> _T: ...
 def sin(x: _T) -> _T: ...
 def sinh(x: _T) -> _T: ...