[tvm] branch main updated: [Relay][Bugfix] fix axis parsing of repeat converter in the MXNet frontend (#15891)

[syfeng]

This is an automated email from the ASF dual-hosted git repository.

syfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 899435956a [Relay][Bugfix] fix axis parsing of repeat converter in the 
MXNet frontend (#15891)
899435956a is described below

commit 899435956a4dc4476eb6e2eb2829d2defc1b341f
Author: Qingchao Shen <qingchaos...@outlook.com>
AuthorDate: Mon Oct 9 22:36:26 2023 +0800

    [Relay][Bugfix] fix axis parsing of repeat converter in the MXNet frontend 
(#15891)
    
    fix axis in repeat
---
 python/tvm/relay/frontend/mxnet.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python/tvm/relay/frontend/mxnet.py 
b/python/tvm/relay/frontend/mxnet.py
index 7ef8c6134e..c3e14c9b72 100644
--- a/python/tvm/relay/frontend/mxnet.py
+++ b/python/tvm/relay/frontend/mxnet.py
@@ -932,7 +932,12 @@ def _mx_repeat(inputs, attrs):
     assert len(inputs) == 1
     new_attrs = {}
     new_attrs["repeats"] = attrs.get_int("repeats")
-    new_attrs["axis"] = attrs.get_int("axis", 0)
+    axis = attrs.get_int("axis", None)
+    if axis is None:
+        inputs[0] = _op.nn.batch_flatten(inputs[0])
+        new_attrs["axis"] = 0
+    else:
+        new_attrs["axis"] = axis
     return _op.repeat(inputs[0], **new_attrs)