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.