AndrewZhaoLuo commented on a change in pull request #9039:
URL: https://github.com/apache/tvm/pull/9039#discussion_r712811152
##########
File path: python/tvm/relay/frontend/onnx.py
##########
@@ -1462,6 +1462,26 @@ def _impl_v1(cls, inputs, attr, params):
inputs[0] = _op.expand_dims(inputs[0], axis=axis, num_newaxis=1)
return inputs[0]
+ @classmethod
+ def _impl_v12(cls, inputs, attr, params):
+ rank_input = len(infer_type(inputs[0]).checked_type.shape)
+ num_new_axis = int(infer_type(inputs[1]).checked_type.shape[0])
+ axes = relay.split(inputs[1], num_new_axis).astuple()
+
+ result = inputs[0]
+
+ # TODO (AndrewZhaoLuo): investigate performance issues with consecutive
+ # dynamic expand_dims on non-llvm targets.
+ for i in range(num_new_axis):
+ axis = relay.TupleGetItem(axes, i)
+ # Unpack scalar
+ axis = relay.reshape(axis, [])
+ axis = relay.If(
+ axis >= relay.const(0, "int64"), axis, axis +
relay.const(rank_input, "int64")
+ )
+ result = _op.expand_dims(result, axis)
+ return result
+
Review comment:
This might be doable if we change the interface to accept a list of
sorted axis but it'll be quite a bit more complicated.
--
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]
