Rebasing

python/mxnet/contrib/onnx/mx2onnx/_op_translations.py

@@ -2059,10 +2059,14 @@ def convert_topk(node, **kwargs):
     axis = int(attrs.get('axis', '-1'))
     k = int(attrs.get('k', '1'))
     ret_type = attrs.get('ret_typ')
+    dtype = attrs.get('dtype')
     outputs = [name + '_output0']
 
     if ret_type and ret_type == 'both':
-        outputs.append(name + '_output1')
+        if dtype and dtype == 'int64':
+            outputs.append(name + '_output1')
+        else:
+            raise NotImplementedError("ONNX expects indices to be of type int64")
     else:
         raise NotImplementedError("ONNX expects both value and indices as output")
 

python/mxnet/contrib/onnx/onnx2mx/_op_translations.py

@@ -784,5 +784,6 @@ def lpnormalization(attrs, inputs, proto_obj):
 def topk(attrs, inputs, proto_obj):
     """Returns the top k elements in an input array along the given axis."""
     new_attrs = translation_utils._add_extra_attributes(attrs,
-                                                        {'ret_typ': 'both'})
+                                                        {'ret_typ': 'both',
+                                                         'dtype': 'int64'})
     return 'topk', new_attrs, inputs