Lunderberg commented on code in PR #16757:
URL: https://github.com/apache/tvm/pull/16757#discussion_r1534339271
##########
python/tvm/relax/frontend/nn/core.py:
##########
@@ -591,7 +591,22 @@ def wrap_nested(expr: rx.Expr, name: str) -> Union[Tensor,
Sequence[Tensor]]:
The computed result.
"""
if not isinstance(expr, rx.DataflowVar):
- expr = BlockBuilder.current().emit(expr, name)
+ block_builder = BlockBuilder.current()
+ if block_builder is None:
+ # Normalize to make sure we have valid StructInfo, but
+ # wait until we are actually building the function to
+ # flatten nested expressions.
+ #
+ # TODO(Lunderberg): Make this easier to call. Infering
+ # struct info for a nested expression should be doable in
+ # a free function, without requiring an active
+ # BlockBuilder and an active FunctionFrame.
+ builder = BlockBuilder()
+ with builder.function("dummy_scope", params=[]):
+ expr = builder.normalize(expr)
+ builder.emit_func_output([])
Review Comment:
The expression isn't actually being generated in the
`builder.function("dummy_scope")` here, because it uses `builder.normalize` and
not `builder.emit`. The purpose is to infer the `StructInfo` to be used, but
to skip the flattening of nested expressions.
The only reason why the function scope is entered is because the
`builder.normalize` method requires an active scope. This is something I hope
to change in the future.
--
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]
