slyubomirsky commented on code in PR #15026:
URL: https://github.com/apache/tvm/pull/15026#discussion_r1224662154
##########
python/tvm/relax/utils.py:
##########
@@ -455,10 +462,14 @@ def _shape_with_old_tir_var(
# with old set of variables.
tir_var_inverse_map = {v: k for k, v in tir_var_map.items()}
- output_sinfo = [
- TensorStructInfo(_shape_with_old_tir_var(out.shape,
tir_var_inverse_map), out.dtype)
- for out in outs
- ]
+ def te_to_sinfo(arg):
+ return TensorStructInfo(_shape_with_old_tir_var(arg.shape,
tir_var_inverse_map), arg.dtype)
+
+ input_sinfo = [te_to_sinfo(arg) for arg in te_args]
+ output_sinfo = [te_to_sinfo(out) for out in outs]
+
+ primfunc_sinfo = FuncStructInfo([*input_sinfo, *output_sinfo],
PrimStructInfo("void"))
+ _update_struct_info(tir_func, primfunc_sinfo)
Review Comment:
@Lunderberg I've put this PR on the agenda for next week's community
meeting. If you can make it, that would be good
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
