apache / apache/tvm

[Relax][Frontend][Torch] `_one_hot` does not validate `num_classes`: non-positive values crash `relax.op.one_hot` with an opaque `InternalError` ("depth must be positive")

Closed
#20,319 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

### Expected behavior

`tvm.relax.frontend.torch._one_hot`
(`python/tvm/relax/frontend/torch/fx_translator.py:725-735` and
`python/tvm/relax/frontend/torch/exported_program_translator.py:1198-1210`) reads the
`num_classes` argument of a PyTorch `F.one_hot` call and forwards it verbatim to
`relax.op.one_hot`. It should validate that `num_classes > 0` and raise a clear,
frontend-level error when it is not — instead of letting an arbitrary value reach the
C++ op and surface as an opaque low-level assertion.

### Actual behavior

`_one_hot` performs **no validation of `num_classes`**. Any non-positive value
(`0`, `-1`, `-2`, …) is passed straight to `relax.op.one_hot`, whose C++ builder
(`src/relax/op/tensor/manipulate.cc:3112`) enforces `TVM_FFI_ICHECK(depth > 0)` and
raises an unhelpful `InternalError`:

```
InternalError: Check failed: (depth > 0) is false: one_hot: depth must be positive, but got 0
```

The failure message gives no hint that the `num_classes` argument is at fault or what
the fix is. Notably, `num_classes=0` — which is a plain constant — is **accepted by both
`torch.export.export` and `fx.symbolic_trace`**, so a data-driven model (e.g. a class
count that evaluates to 0) converts through the torch toolchain without complaint and
then crashes TVM:

```
# torch.export succeeds:
ep = torch.export.export(lambda x: F.one_hot(x, num_classes=0), (x,)) # OK
# fx.symbolic_trace succeeds:
gm = fx.symbolic_trace(M) # OK
# ...but TVM's legacy from_fx path crashes on both:
mod = from_fx(gm, [((3,), "int64")])
# InternalError: Check failed: (depth > 0) is false: one_hot: depth must be positive, but got 0
```

Additional context:

- In the **modern path** (`from_exported_program`), `_one_hot` is effectively **dead
code**: `from_exported_program` runs `exported_program.run_decompositions()` by
default (`exported_program_translator.py:2386-2387`), which rewrites
`aten.one_hot` into `arange + unsqueeze + eq + _to_copy` before dispatch — so the
recommended path never reaches `relax.op.one_hot` (verified by instrumentation,
`_one_hot` call count = 0). The validation gap is therefore only triggerable via the
legacy `from_fx` path, but it lives in shared converter logic and the recommended
path itself relies on torch's decomposition for every `num_classes` value.
- The documented PyTorch default `num_classes=-1` (auto-infer `max+1`) is a
data-dependent value and is rejected by `torch.export` itself
(`GuardOnDataDependentSymNode`) before TVM is reached; the legacy `from_fx` path on
the same default reports the misleading
`ValueError: num_classes not found in node.args or node.kwargs` (the argument is
absent rather than `-1`).

### Environment

- OS: Linux
- TVM: v0.24.dev0 (main branch, commit `390af87345`, built 2026-08-18)
- Python: 3.11
- torch: 2.10.0+cu128
- target: `llvm`

### Steps to reproduce

```python
"""Repro: TVM torch frontend crashes on non-positive num_classes (no validation)."""
import torch
import torch.nn.functional as F
from tvm import relax
from tvm.relax.frontend.torch import from_fx

x = torch.tensor([0, 2, 1], dtype=torch.int64)

class M(torch.nn.Module):
def forward(self, x):
return F.one_hot(x, num_classes=0) # non-positive constant

# 1) The torch toolchain accepts the model:
import torch.fx as fx
gm = fx.symbolic_trace(M()) # OK
print("fx.symbolic_trace:", " ".join(str(n) for n in gm.graph.nodes))

# 2) TVM legacy from_fx path crashes at conversion:
try:
mod = from_fx(gm, [((3,), "int64")])
ex = relax.build(mod, target="llvm")
print("TVM from_fx: OK")
except Exception as e:
print(f"TVM from_fx: {type(e).__name__}: {e}")

# 3) torch.export also accepts num_classes=0 (context):
ep = torch.export.export(M(), (x,))
print("torch.export:", " ".join(str(n.target) for n in ep.graph.nodes if n.op == "call_function"))
```

Actual output:

```
fx.symbolic_trace: x one_hot output
TVM from_fx: InternalError: Check failed: (depth > 0) is false: one_hot: depth must be positive, but got 0
torch.export: one_hot.default
```

The same crash occurs for `num_classes=-1` and `num_classes=-2`
(`... depth must be positive, but got -1` / `-2`).

### Triage

* needs-triage
* bug
* relax
* frontend/torch

Contributor guide

No contributing guide indexed for this repository

Research direction

Start at _one_hot in python/tvm/relax/frontend/torch/fx_translator.py:725-735 and exported_program_translator.py:1198-1210, then reproduce the failure through the legacy from_fx entry point with non-positive num_classes. Trace how the value reaches relax.op.one_hot and ensure invalid values produce a clear frontend-level error; done means the opaque InternalError is replaced and the reported cases are covered by tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.