Update Argmin/Argmax ONNX Export (#38329)

Summary:
Update Argmin/Argmax ONNX export in opset 12 to export with "select_last_index", and export correctly cases where the same value appears multiple time in the input tensor.
Pull Request resolved: #38329

Reviewed By: hl475

Differential Revision: D21613799

Pulled By: houseroad

fbshipit-source-id: 4597e23561f444c4e56d30c735dae7e9a8a41c5e

Author: lara-hdr

test/onnx/test_pytorch_onnx_onnxruntime.py

@@ -2045,6 +2045,32 @@ def forward(self, input, other):
         y = torch.randint(10, (2, 4, 5))
         self.run_test(MatmulModel(), (x, y))
 
+    def _argmin_argmax_model(self, input):
+        class ArgminArgmaxModel(torch.nn.Module):
+            def forward(self, input):
+                return torch.argmin(input), \
+                    torch.argmax(input), \
+                    torch.argmin(input, keepdim=True), \
+                    torch.argmax(input, keepdim=True)
+
+        self.run_test(ArgminArgmaxModel(), input)
+
+    def test_argmin_argmax(self):
+        input = torch.randn(7, 3, 5)
+        self._argmin_argmax_model(input)
+
+    # Argmin and Argmax with "select_last_index" is not supprted before opset 12
+    # "select_last_index" was added in opset 12 to deal with corner case where the
+    # same value appears multiple times in the tensor
+    @skipIfUnsupportedMinOpsetVersion(12)
+    def test_argmin_argmax_select_last_index(self):
+        input = torch.tensor([[1., 2., 3.],
+                             [1., 1., 2.]])
+        self._argmin_argmax_model(input)
+
+        input = torch.ones(7, 3, 5)
+        self._argmin_argmax_model(input)
+
     def test_view(self):
         class ViewModel(torch.nn.Module):
             def forward(self, input):

torch/onnx/symbolic_opset12.py

@@ -2,7 +2,7 @@
 
 import torch
 import torch.onnx.symbolic_helper as sym_help
-from torch.onnx.symbolic_helper import parse_args
+from torch.onnx.symbolic_helper import parse_args, _parse_arg
 
 
 # EDITING THIS FILE? READ THIS FIRST!
@@ -62,6 +62,28 @@ def nll_loss2d(g, self, target, weight, reduction, ignore_index):
     return nll_loss(g, self, target, weight, reduction, ignore_index)
 
 
+def argmax(g, input, dim, keepdim):
+    if sym_help._is_none(dim):
+        from torch.onnx.symbolic_opset9 import reshape
+        flattened = reshape(g, input, (-1,))
+        return g.op('ArgMax', flattened, axis_i=0, keepdims_i=False, select_last_index_i=True)
+    else:
+        dim = _parse_arg(dim, 'i')
+        keepdim = _parse_arg(keepdim, 'i')
+        return g.op('ArgMax', input, axis_i=dim, keepdims_i=keepdim, select_last_index_i=True)
+
+
+def argmin(g, input, dim, keepdim):
+    if sym_help._is_none(dim):
+        from torch.onnx.symbolic_opset9 import reshape
+        flattened = reshape(g, input, (-1,))
+        return g.op('ArgMin', flattened, axis_i=0, keepdims_i=False, select_last_index_i=True)
+    else:
+        dim = _parse_arg(dim, 'i')
+        keepdim = _parse_arg(keepdim, 'i')
+        return g.op('ArgMin', input, axis_i=dim, keepdims_i=keepdim, select_last_index_i=True)
+
+
 def pow(g, self, exponent):
     return g.op("Pow", self, exponent)