This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 92125d79cf [BugFix][Unity] Fix for TupleStructInfo (#15538)
92125d79cf is described below
commit 92125d79cfc2a93b07ab077e468fe30c77a22a4b
Author: Xiyou Zhou <[email protected]>
AuthorDate: Sat Aug 12 13:23:34 2023 -0700
[BugFix][Unity] Fix for TupleStructInfo (#15538)
Fix for TupleStructInfo.
---
python/tvm/dlight/benchmark/extract.py | 6 ++++++
1 file changed, 6 insertions(+)
diff --git a/python/tvm/dlight/benchmark/extract.py
b/python/tvm/dlight/benchmark/extract.py
index df9b06e3fd..ddb9eb58ff 100644
--- a/python/tvm/dlight/benchmark/extract.py
+++ b/python/tvm/dlight/benchmark/extract.py
@@ -119,7 +119,13 @@ def extract_dynamic_var(
dym_var_dict[gv] = {}
for functor in func_dict[gv]:
for arg_list, _ in func_dict[gv][functor]:
+ flattened_arg_list = []
for arg in arg_list:
+ if isinstance(arg, relax.TupleStructInfo):
+ flattened_arg_list.extend(arg.fields)
+ else:
+ flattened_arg_list.append(arg)
+ for arg in flattened_arg_list:
if isinstance(arg, relax.TensorStructInfo):
for val in arg.shape.values:
if isinstance(val, tvm.tir.Var):
