Fix aten_amax and aten_amin export when dim is omitted

[om singhal (Om-singhaI)]

torch.amax(x) and torch.amin(x) don't export. The aten schema is amax(Tensor self, int[1] dim=[], bool keepdim=False), so dim defaults to the empty list, meaning reduce every dimension. aten_amax and aten_amin in onnxscript/function_libs/torch_lib/ops/core.py declared dim with no default, which makes it required, and torch.export emits aten.amax.default(x) with no dim at all when the caller leaves it out.

class Model(torch.nn.Module):
    def forward(self, x):
        return torch.amax(x)

torch.onnx.export(Model(), (torch.randn(2, 3),), dynamo=True)

ValueError: Required parameter 'dim' is not provided. Signature:
pkg.onnxscript.torch_lib::aten_amax(self: TRealOrUInt8, dim: T_dim, keepdim: INT = 0) -> (TRealOrUInt8)
... Args: (SymbolicTensor(name='x', type=Tensor(FLOAT), shape=Shape([2, 3])),). Kwargs: {}.

The same reduction exports fine one line over

torch.amax(x, keepdim=True) works on main today. torch.export can't reach keepdim without filling dim in positionally, so it emits aten.amax.default(x, [], True) and the lowering handles that empty axes input correctly. Same op, same empty dim, one spelling exports and the other raises. The only difference is whether dim reached the graph at all.

The comment already sitting in the source says ReduceMax reduces all dimensions when dim is empty, so the intent was right. The signature just never let you get there.

The fix

aten_amax and aten_amin become trace_only and default dim to None. On that path they call ReduceMax and ReduceMin with no axes input.

That last part is the bit worth checking. The ONNX spec for ReduceMax 18 says an empty axes, "either not provided or explicitly empty", reduces over all axes when noop_with_empty_axes is false and over the empty set of axes when it's true. It defaults to 0, so leaving axes out reduces everything. A 1 there would hand back the input untouched, which is a wrong answer rather than an error, so the test asserts the attribute is 0 on the emitted node as well as comparing against eager.

I ran both forms against onnxruntime, opset 18, on a (2, 3) input whose max is 9.0:

axes input ABSENT:          shape=() value=9.0
explicit EMPTY axes tensor: shape=() value=9.0

The graph for torch.amax(x) with this change is one node:

ReduceMax in=['x'] attrs={'noop_with_empty_axes': 0, 'keepdims': 0}

Tests

test_amax_amin_reduce_every_dimension_when_dim_is_omitted in tests/function_libs/torch_lib/e2e_ops_tests.py, four cases: amax and amin, each with and without keepdim. The keepdim pair passes on main and is there to pin that both spellings keep agreeing.

pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k amax_amin:

main, new test only:  2 failed, 2 passed
with this change:     4 passed

pytest tests/function_libs/torch_lib/ops_test.py -k "amax or amin":

main:               18 passed, 4 skipped, 328 subtests passed
with this change:   16 passed, 6 skipped, 328 subtests passed

The two extra skips are the function proto validity checks, which skip for traced functions. Same 328 subtests pass either way.

Whole file, pytest tests/function_libs/torch_lib/e2e_ops_tests.py: 8 failed, 119 passed, 1 skipped, 104 subtests passed. The same 8 fail on main without my change (stft, deform_conv2d, sdpa bool mask, unbind dynamic, convolution complex kernel shape), so they're unrelated.

ruff check and ruff format --check pass on both files with ruff 0.15.1, the lintrunner pinned version.

Environment: Python 3.10, torch 2.9.1, onnx 1.22.0, onnxruntime 1.23.2, macOS arm64.

[Copilot]

🟢 Approval recommended

The change is narrowly scoped, matches ONNX Reduce* semantics for omitted axes, and is covered by a targeted end-to-end regression test.

Pull request overview

This PR fixes ONNX export for torch.amax(x) / torch.amin(x) when dim is omitted by aligning the aten::amax / aten::amin lowering signatures with the ATen schema default (dim=[] meaning reduce over all dimensions), and ensuring the lowering can handle the “argument omitted entirely” form emitted by torch.export.

Changes:

• Make aten_amax / aten_amin accept an omitted dim (defaulting to None) and lower to ReduceMax / ReduceMin without providing an axes input in that case.
• Add an end-to-end regression test covering amax/amin with and without keepdim when dim is omitted.
• Add a graph-level check intended to guard against noop_with_empty_axes=1 on the emitted reduce nodes.

File summaries

File	Description
tests/function_libs/torch_lib/e2e_ops_tests.py	Adds an E2E regression test for amax/amin export when dim is omitted, including a guard against noop_with_empty_axes being set incorrectly.
onnxscript/function_libs/torch_lib/ops/core.py	Updates aten::amax / aten::amin signatures and lowering to correctly handle omitted dim by reducing over all axes.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.69%. Comparing base (d1c005d) to head (a6e1cc5).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	50.00%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3035      +/-   ##
==========================================
- Coverage   72.70%   72.69%   -0.01%     
==========================================
  Files         265      265              
  Lines       32298    32302       +4     
  Branches     3059     3061       +2     
==========================================
+ Hits        23481    23483       +2     
  Misses       7779     7779              
- Partials     1038     1040       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.