This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new fd7c81de3b [TIR][Schedule] Remove `@type_check` for
`set_axis_separator` (#17134)
fd7c81de3b is described below
commit fd7c81de3b2f1e023351b10478e647fdbf367acc
Author: Abhikrant Sharma <[email protected]>
AuthorDate: Tue Jul 9 22:25:38 2024 +0530
[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.
---
python/tvm/relax/transform/legalize_ops/manipulate.py | 1 -
python/tvm/tir/schedule/schedule.py | 1 -
2 files changed, 2 deletions(-)
diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py
b/python/tvm/relax/transform/legalize_ops/manipulate.py
index 4d30b97f64..1efa78c069 100644
--- a/python/tvm/relax/transform/legalize_ops/manipulate.py
+++ b/python/tvm/relax/transform/legalize_ops/manipulate.py
@@ -213,7 +213,6 @@ def _layout_transform(bb: BlockBuilder, call: Call) -> Expr:
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
index f477a0f112..4127266da7 100644
--- a/python/tvm/tir/schedule/schedule.py
+++ b/python/tvm/tir/schedule/schedule.py
@@ -3490,7 +3490,6 @@ class Schedule(Object):
self, block, index_map
)
- @type_checked
def set_axis_separator(
self,
block: Union[BlockRV, str],
