hiyufan opened a new pull request, #20254:
URL: https://github.com/apache/tvm/pull/20254
### 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:
```python
# _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:
```python
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
```python
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
