Lunderberg commented on code in PR #16757:
URL: https://github.com/apache/tvm/pull/16757#discussion_r1536941583
##########
python/tvm/relax/frontend/nn/exporter.py:
##########
@@ -176,35 +183,26 @@ def _unwrap_ret(expr: typing.Any) -> typing.Any:
def _convert_input(arg):
if isinstance(arg, tir.Var):
return rx.Var(arg.name, struct_info=ShapeStructInfo(values=[arg]))
- if isinstance(arg, (core.Tensor, core.Object)):
+ elif isinstance(arg, (core.Tensor, core.Object)):
return arg._expr # pylint: disable=protected-access
- if isinstance(arg, _spec.Tuple):
+ elif isinstance(arg, _spec.Tuple):
return rx.Var(
arg.name,
struct_info=TupleStructInfo(
[_convert_input(arg_i).struct_info for arg_i in
arg.elements]
),
)
- raise TypeError(f"Unsupported input type: {type(arg)}")
+ elif isinstance(arg, rx.Expr):
+ return arg
+ else:
+ raise TypeError(f"Unsupported input type: {type(arg)}")
def _params(mode: str) -> typing.List[rx.Var]:
inputs: typing.List[rx.Var] = []
- def _get_var(shape_var: tir.Var) -> tir.Var:
- name = shape_var.name
- if name in str2var_params:
- return str2var_params[name]
Review Comment:
@MasterJH5574 Thank you, and that helps to reproduce the issue. Does the
issue only involve `nn.Linear`, or would it also apply to dynamic shapes that
span across multiple independent `nn.Module` subclasses? If the issue is
specific to `nn.Linear`, can you take a look at
https://github.com/apache/tvm/pull/16781? This avoids the duplicate default
conversion from `str` to a fresh TIR variable in both the weights and bias.
Overall, I think de-duplicating `str` values makes sense, because Python
`str` objects have value equality. De-duplicating TIR variables by their name
doesn't make sense, because TIR variables have reference equality. I think we
could have a good intermediate implementation if we de-duplicate `str` values
prior to conversion into TIR variables.
--
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]
