Lunderberg commented on code in PR #15026:
URL: https://github.com/apache/tvm/pull/15026#discussion_r1226872667
##########
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:
Thank you! I probably won't be able to attend, given the timing, but I
agree that discussion would be good. For now, I've converted this PR to a
draft, to ensure that it can't be merged prior to discussion.
--
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]
