Fix aten_amax and aten_amin export when dim is omitted

onnxscript/function_libs/torch_lib/ops/core.py

@@ -419,19 +419,33 @@ def aten_alpha_dropout(input: TensorType, p: float, train: bool) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::amax")
-def aten_amax(self: TRealOrUInt8, dim: INT64, keepdim: bool = False) -> TRealOrUInt8:
+@torch_op("aten::amax", trace_only=True)
+def aten_amax(
+    self: TRealOrUInt8, dim: Optional[INT64] = None, keepdim: bool = False
+) -> TRealOrUInt8:
     """amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor"""
 
-    # ReduceMax reduces all dimensions when dim is empty
+    if dim is None:
+        # dim defaults to the empty list in the aten schema, which means reduce every
+        # dimension. noop_with_empty_axes keeps its default of 0, so ReduceMax without
+        # an axes input reduces all of them.
+        return op.ReduceMax(self, keepdims=keepdim)
+    # An explicitly empty dim arrives here and reduces every dimension for the same reason
     return op.ReduceMax(self, dim, keepdims=keepdim)
 
 
-@torch_op("aten::amin")
-def aten_amin(self: TRealOrUInt8, dim: INT64, keepdim: bool = False) -> TRealOrUInt8:
+@torch_op("aten::amin", trace_only=True)
+def aten_amin(
+    self: TRealOrUInt8, dim: Optional[INT64] = None, keepdim: bool = False
+) -> TRealOrUInt8:
     """amin(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor"""
 
-    # ReduceMin reduces all dimensions when dim is empty
+    if dim is None:
+        # dim defaults to the empty list in the aten schema, which means reduce every
+        # dimension. noop_with_empty_axes keeps its default of 0, so ReduceMin without
+        # an axes input reduces all of them.
+        return op.ReduceMin(self, keepdims=keepdim)
+    # An explicitly empty dim arrives here and reduces every dimension for the same reason
     return op.ReduceMin(self, dim, keepdims=keepdim)
 
 

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -1788,6 +1788,37 @@ def forward(self, a, b):
         onnx_program = torch.onnx.export(IsCloseModel(), (a, b), dynamo=True, optimize=False)
         _testing.assert_onnx_program(onnx_program)
 
+    @parameterized.parameterized.expand(
+        [
+            ("amax", "amax", False),
+            ("amax_keepdim", "amax", True),
+            ("amin", "amin", False),
+            ("amin_keepdim", "amin", True),
+        ]
+    )
+    def test_amax_amin_reduce_every_dimension_when_dim_is_omitted(
+        self, _: str, reduction: str, keepdim: bool
+    ):
+        # dim defaults to the empty list in the aten schema, so leaving it out means
+        # reduce every dimension. torch.export drops the argument entirely unless a
+        # later one is set, in which case it passes an empty list instead, and both
+        # spellings have to come out the same. ReduceMax and ReduceMin only reduce
+        # everything while noop_with_empty_axes is 0. A 1 there would quietly hand
+        # back the input untouched.
+        reduce_op = getattr(torch, reduction)
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return reduce_op(x, keepdim=keepdim)
+
+        onnx_program = torch.onnx.export(
+            Model(), (torch.randn(2, 3),), dynamo=True, optimize=False
+        )
+        for node in onnx_program.model.graph:
+            if node.op_type in ("ReduceMax", "ReduceMin"):
+                self.assertEqual(node.attributes.get_int("noop_with_empty_axes", 0), 0)
+        _testing.assert_onnx_program(onnx_program)
+
 
 if __name__ == "__main__":
     unittest.main()