[torchlib] Register aten::sort.stable and implement aten::argsort

[Anish Mehta (anishmehta24)]

torch.sort(x, stable=True) and torch.argsort(x, stable=True) lower to the aten::sort.stable / aten::argsort.stable overloads, and torch.argsort(x) to aten::argsort. None of them had an ONNX function: aten_sort was registered for aten::sort only, and aten_argsort was a NotImplementedError stub, so torch.onnx.export(..., dynamo=True) failed with

DispatchError: No ONNX function found for <OpOverload(op='aten.sort', overload='stable')>. Failure message: No decompositions registered for the real-valued input

Repro:

class M(torch.nn.Module):
    def forward(self, a):
        return torch.sort(a, dim=1, descending=True, stable=True).values
torch.onnx.export(M(), (torch.randn(5, 3, 4),), dynamo=True)

Changes

• aten_sort now also registers aten::sort.stable and accepts the overload's bool? stable keyword. TopK is the only ONNX op that sorts and gives no stability guarantee, so the flag is accepted and ignored, exactly as it already was for aten::sort; both overloads lower to the same graph.
• aten_argsort is implemented on top of aten_sort (returning the indices) and registered for aten::argsort and aten::argsort.stable.
• TorchLibOpInfo("argsort", ...) added to the op tests, with the same float16 xfail as sort: the 8000-element float16 sample has many exact ties whose relative order is unspecified (float32 passes).

Validation

• pytest tests/function_libs/torch_lib/ops_test.py -k sort: 6 passed, 2 xfailed (float16 sort/argsort), 348 subtests passed.
• End-to-end with torch 2.14 + onnxruntime 1.30: sort(stable=True) values+indices, descending, argsort default / descending / stable=True all export and match torch eager, including on inputs with duplicate values.
• ruff check / ruff format clean on the touched files.

AI assistance: drafted with Claude Code and validated as above.

🤖 Generated with Claude Code

https://claude.ai/code/session_01D56yuT4mLP36Um1WF51fEA