[Fix][Relax][Frontend][Torch] Emit int64 indices for sort and argsort

[hiyufan]

Problem

torch.sort and torch.argsort return int64 indices. The Relax PyTorch frontend emits int32 for them.

relax.op.argsort defaults to dtype="int32", and both call sites in base_fx_graph_translator.py take that default:

# _argsort
return self.block_builder.emit(relax.op.argsort(x, dim, descending))

# _sort
indices = self.block_builder.emit(relax.op.argsort(x, dim, descending))

To be precise about the scope: the sorted values and the index values are already correct. Only the index dtype diverges from PyTorch.

Why this looks like an oversight rather than a deliberate choice

relax.op.topk has the same dtype="int32" default, and _topk in this very file explicitly overrides it:

relax.op.topk(x, k=k, axis=dim, largest=largest, ret_type="both", dtype="int64")

So the frontend already establishes that the relax-level default has to be overridden to match PyTorch. Every index-producing torch op the frontend supports comes out as int64 — except sort and argsort:

torch op	torch index dtype	frontend before this PR
torch.sort	int64	int32
torch.argsort	int64	int32
torch.topk	int64	int64 (explicit dtype=)
torch.argmax / torch.argmin	int64	int64
torch.max(dim=) / torch.median(dim=)	int64	int64
torch.bucketize	int64	int64

A single imported graph that uses both topk and sort therefore carries two different index dtypes for the same kind of value.

Reproduce

import torch
from torch import nn
from tvm.relax.frontend.torch import from_exported_program


def relax_dtypes(fn, x):
    class M(nn.Module):
        def forward(self, t):
            return fn(t)

    ep = torch.export.export(M().eval(), (x,))
    return [str(f.dtype) for f in from_exported_program(ep)["main"].ret_ty.fields]


x = torch.randn(3, 4)
print("sort   ", relax_dtypes(lambda t: torch.sort(t, dim=1), x))
print("argsort", relax_dtypes(lambda t: torch.argsort(t, dim=1), x))
print("topk   ", relax_dtypes(lambda t: torch.topk(t, 2, dim=1), x))

Before / after:

                 before                        after
sort     ['float32', 'int32']          ['float32', 'int64']
argsort  ['int32']                     ['int64']
topk     ['float32', 'int64']          ['float32', 'int64']

Fix

Pass dtype="int64" at both relax.op.argsort call sites, matching what _topk already does.

Verification

Built with LLVM and ran the imported graph through the VM against PyTorch on sort(dim=1), sort(descending=True), sort(dim=0) and argsort(dim=1):

• before: values equal, indices equal, index dtype int32 vs torch int64
• after: values equal, indices equal, index dtype int64 — matches torch on every case

Test suites, tests/python/relax/test_frontend_from_fx.py + tests/python/relax/test_frontend_from_exported_program.py:

• clean main: 24 failed, 412 passed, 3 skipped
• with this change: 24 failed, 413 passed, 3 skipped

The 24 failures are pre-existing on main in my environment (test_dtypes and friends) and are identical before and after; the one additional pass is the new test_sort.

ruff format --check and ruff check are clean on the touched files.

Tests

• Updated the expected IR in test_argsort (both frontend test files) and test_sort (test_frontend_from_fx.py) — these had the int32 result written into them.
• Added test_sort to test_frontend_from_exported_program.py; that path had no coverage for torch.sort.

---

This change was prepared with AI assistance (Claude). I have reviewed and verified it, and can speak to it in review.