Fix docs.

Author: zxybazh

python/tvm/meta_schedule/database/database.py

@@ -284,7 +284,7 @@ def query_schedule(
 
         Returns
         -------
-        schedule : Optional[Schedule]
+        schedule : Optional[tvm.tir.Schedule]
             The best schedule of the given workload; None if not found.
         """
         return _ffi_api.DatabaseQuerySchedule(self, mod, target, workload_name)  # type: ignore # pylint: disable=no-member
@@ -338,7 +338,7 @@ def query(
 
         Returns
         -------
-        result : Union[Schedule, IRModule, TuningRecord]
+        result : Union[tvm.tir.Schedule, IRModule, TuningRecord]
             The best optimization outcome of the given workload.
         """
         if kind == "schedule":
@@ -381,7 +381,8 @@ def create(  # pylint: disable=keyword-arg-before-vararg
 
         Parameters
         ----------
-        kind : str = "json" | "memory" | "union" | "ordered_union" | Callable[[Schedule], bool]
+        kind : str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.tir.Schedule],
+        bool]
             The kind of the database to be created. The following kinds are supported:
             "json", "memory", "union", "ordered_union", and a custom schedule function.
 

python/tvm/meta_schedule/postproc/postproc.py

@@ -54,7 +54,7 @@ def apply(self, sch: Schedule) -> bool:
 
         Parameters
         ----------
-        sch : Schedule
+        sch : tvm.tir.Schedule
             The schedule to be post processed.
 
         Returns

python/tvm/meta_schedule/schedule_rule/schedule_rule.py

@@ -57,14 +57,14 @@ def apply(self, sch: Schedule, block: BlockRV) -> List[Schedule]:
 
         Parameters
         ----------
-        sch : Schedule
+        sch : tvm.tir.Schedule
             The schedule to be modified.
         block : BlockRV
             The specific block to apply the schedule rule.
 
         Returns
         -------
-        design_spaces : List[Schedule]
+        design_spaces : List[tvm.tir.Schedule]
             The list of schedules generated by applying the schedule rule.
         """
         return _ffi_api.ScheduleRuleApply(  #  type: ignore # pylint: disable=no-member

python/tvm/meta_schedule/search_strategy/search_strategy.py

@@ -44,7 +44,7 @@ class MeasureCandidate(Object):
 
     Parameters
     ----------
-    sch : Schedule
+    sch : tvm.tir.Schedule
         The schedule to be measured.
     args_info : List[ArgInfo]
         The argument information.
@@ -62,7 +62,7 @@ def __init__(
 
         Parameters
         ----------
-        sch : Schedule
+        sch : tvm.tir.Schedule
             The schedule to be measured.
         args_info : List[ArgInfo]
             The argument information.
@@ -115,7 +115,7 @@ def pre_tuning(
             The maximum number of trials.
         num_trials_per_iter : int
             The number of trials per iteration.
-        design_spaces : List[Schedule]
+        design_spaces : List[tvm.tir.Schedule]
             The design spaces used during tuning process.
         database : Optional[Database] = None
             The database used during tuning process.

python/tvm/meta_schedule/space_generator/space_generator.py

@@ -80,7 +80,7 @@ def generate_design_space(self, mod: IRModule) -> List[Schedule]:
 
         Returns
         -------
-        design_spaces : List[Schedule]
+        design_spaces : List[tvm.tir.Schedule]
             The generated design spaces, i.e., schedules.
         """
         return _ffi_api.SpaceGeneratorGenerateDesignSpace(self, mod)  # type: ignore # pylint: disable=no-member
@@ -248,7 +248,7 @@ def generate_design_space(self, mod: IRModule) -> List[Schedule]:
 
         Returns
         -------
-        design_spaces : List[Schedule]
+        design_spaces : List[tvm.tir.Schedule]
             The generated design spaces, i.e., schedules.
         """
         raise NotImplementedError

python/tvm/meta_schedule/tune.py

@@ -73,13 +73,13 @@ def tune_tasks(
     module_equality : Optional[str]
         A string to specify the module equality testing and hashing method.
         It must be one of the followings:
-          - "structural": Use StructuralEqual/Hash
-          - "ignore-ndarray": Same as "structural", but ignore ndarray raw data during
-                              equality testing and hashing.
-          - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
-                            given module. The "ignore-ndarray" varint is used for the extracted
-                            blocks or in case no anchor block is found.
-                            For the definition of the anchor block, see tir/analysis/analysis.py.
+            - "structural": Use StructuralEqual/Hash
+            - "ignore-ndarray": Same as "structural", but ignore ndarray raw data during equality
+                testing and hashing.
+            - "anchor-block": Apply equality testing and hashing on the anchor block extracted from
+                a given module. The "ignore-ndarray" varint is used for the extracted blocks or in
+                case no anchor block is found. For the definition of the anchor block, see
+                tir/analysis/analysis.py.
 
     Returns
     -------

python/tvm/meta_schedule/tune_context.py

@@ -154,7 +154,7 @@ def generate_design_space(self) -> List[Schedule]:
 
         Returns
         -------
-        design_spaces : List[Schedule]
+        design_spaces : List[tvm.tir.Schedule]
             The generated design spaces, i.e., schedules.
         """
         if self.mod is None:
@@ -184,7 +184,7 @@ def pre_tuning(
             The maximum number of trials to be executed.
         num_trials_per_iter : int = 64
             The number of trials to be executed per iteration.
-        design_spaces : Optional[List[Schedule]]
+        design_spaces : Optional[List[tvm.tir.Schedule]]
             The design spaces used during tuning process.
             If None, use the outcome of `self.generate_design_space()`.
         database : Optional[Database] = None