From a6e1cc51695be768fb60d95e89bc6930fbd33a70 Mon Sep 17 00:00:00 2001 From: Om Singhal Date: Mon, 7 Sep 2026 17:27:30 -0400 Subject: [PATCH] Fix aten_amax and aten_amin export when dim is omitted The aten schema is amax(Tensor self, int[1] dim=[], bool keepdim=False), so dim defaults to the empty list and means reduce every dimension. torchlib declared dim with no default, which made it required, and torch.export emits aten.amax.default(x) with no dim when the caller leaves it out. The dispatcher then raised ValueError: Required parameter 'dim' is not provided. torch.amax(x, keepdim=True) already worked, because torch.export has to materialize dim=[] positionally to reach keepdim, and the lowering handles an empty axes input correctly. Only the spelling that drops dim failed. aten_amax and aten_amin are now trace_only and default dim to None, which takes ReduceMax and ReduceMin without an axes input. noop_with_empty_axes keeps its default of 0, so that reduces every axis rather than acting as an identity. --- .../function_libs/torch_lib/ops/core.py | 26 ++++++++++++---- .../function_libs/torch_lib/e2e_ops_tests.py | 31 +++++++++++++++++++ 2 files changed, 51 insertions(+), 6 deletions(-) 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()