AndrewZhaoLuo commented on a change in pull request #9039:
URL: https://github.com/apache/tvm/pull/9039#discussion_r712809462
##########
File path: python/tvm/relay/op/transform.py
##########
@@ -110,7 +110,17 @@ def expand_dims(data, axis, num_newaxis=1):
result : relay.Expr
The reshaped result.
"""
- return _make.expand_dims(data, axis, num_newaxis)
+ if isinstance(axis, int):
+ return _make.expand_dims(data, axis, num_newaxis)
+ if isinstance(axis, Expr):
+ # TODO (AndrewZhaoLuo): investigate performance issues with consecutive
+ # dynamic expand_dims on non-llvm targets.
+ for _ in range(num_newaxis):
+ # Dynamic rank is not well supported so we can only increase rank
+ # by a static amount (e.g. 1) so we have to do this
+ data = _dyn_make.expand_dims(data, axis)
+ return data
+ raise ValueError(f"Unknown type for axis: {type(axis)}")
Review comment:
We can, done.
--
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]
