[TIR][Schedule] Remove @type_check for set_axis_separator (#17134)

[TIR][Schedule] Remove @type_check decorator for set_axis_separator The decorator is not allowing types like Array to be passed to set_axis_separator directive. The FFI has a type checking implemented internally making this redundant.
apache · Jul 9, 2024 · fd7c81d · fd7c81d
1 parent c4e6f96
commit fd7c81d
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 2 deletions.
diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py b/python/tvm/relax/transform/legalize_ops/manipulate.py
@@ -213,7 +213,6 @@ def set_axis_sep(axis_sep: list, sch: tir.schedule, buffer_type: str):
     sch.transform_layout(primfunc_name, ("write", 0), index_map, pad_value)
     set_axis_sep(axis_separators, sch, "write")
     if input_axis_separators is not None:
-        input_axis_separators = [int(sep) for sep in input_axis_separators]
         set_axis_sep(input_axis_separators, sch, "read")
     gvar = bb.add_func(sch.mod["main"], primfunc_name)
     output_shape = index_map.map_shape(list(call_args[0].struct_info.shape))

diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py
@@ -3490,7 +3490,6 @@ def after_transform_block_layout(
             self, block, index_map
         )
 
-    @type_checked
     def set_axis_separator(
         self,
         block: Union[BlockRV, str],