diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 27d30e377f..0322f4928e 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/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) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 8a50b5d58d..dc284ebf5c 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/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()