fix comments

Author: yongwww

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -47,8 +47,8 @@ def _get_target(self, expr: Expr) -> Target:
         target = Target.current()
         if target is None:
             raise ValueError(
-                "Target not found. Please ensure that the target is annotated within the module, \
-                or alternatively, execute this within a specified target context."
+                "Target not found. Please ensure that the target is annotated within the module, "
+                "or alternatively, execute this within a specified target context."
             )
         return target
 
@@ -65,12 +65,14 @@ def visit_call_(self, call: Call) -> Expr:
                         call.args[0],
                         call.attrs.axis,
                         not call.attrs.descending,
+                        primfunc_attrs={"tir.is_scheduled": 1},
                     )
                 return self.builder_.call_te(
                     topi.cuda.sort if tgt.kind.name == "cuda" else topi.sort,
                     call.args[0],
                     call.attrs.axis,
                     not call.attrs.descending,
+                    primfunc_attrs={"tir.is_scheduled": 1},
                 )
 
         if call.op.name == "relax.cumsum":
@@ -82,6 +84,7 @@ def visit_call_(self, call: Call) -> Expr:
                     call.args[0],
                     axis,
                     call.attrs.dtype,
+                    primfunc_attrs={"tir.is_scheduled": 1},
                 )
 
         return super().visit_call_(call)

python/tvm/relax/op/sort.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """Sortings operators."""
-from typing import List, Optional
 
 from . import _ffi_api
 from ..expr import Expr
@@ -30,11 +29,11 @@ def sort(x: Expr, axis: int = -1, descending: bool = False):
     x : relax.Expr
         The input tensor.
 
-    axis : Optional[int]
+    axis : int
         Axis along which to sort the input tensor.
         By default the last axis of the input is used.
 
-    descending : Optional[bool]
+    descending : bool
         Whether to sort in descending order, the default is False
 
     Returns