This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a029ef2a77 [Fix][Relax][Frontend][Torch] Validate `num_classes` in the
`one_hot` converters (#20320)
a029ef2a77 is described below
commit a029ef2a77713410b4c97dbc3c527910299b8179
Author: HuEnwei <[email protected]>
AuthorDate: Wed Sep 16 08:07:34 2026 +0800
[Fix][Relax][Frontend][Torch] Validate `num_classes` in the `one_hot`
converters (#20320)
Fixes: #20319
## Summary
The Relax PyTorch frontend's `_one_hot` converter reads the
`num_classes`
argument of an `F.one_hot` / `aten.one_hot` call and forwards it
**verbatim** to
`relax.op.one_hot`. It performs no validation at all, so a non-positive
`num_classes` reaches the C++ op builder
(`src/relax/op/tensor/manipulate.cc`), which asserts:
```
InternalError: Check failed: (depth > 0) is false:
one_hot: depth must be positive, but got 0
```
The message never mentions `num_classes` and gives no hint about how to
fix the
model. This is reachable because `num_classes` is an ordinary constant:
**both `torch.export.export` and `fx.symbolic_trace` accept it as-is**
and record
it in the graph, so the failure only appears once the graph is lowered
through
TVM.
This PR makes the frontend reject a non-positive `num_classes` with a
clear
`ValueError` that names the argument, on **both** converter copies.
## Root cause
`_one_hot` exists twice — once for the legacy `from_fx` path
(`fx_translator.py`) and once for the modern `from_exported_program`
path
(`exported_program_translator.py`). Both validate only that
`num_classes` was
found, then pass it straight through:
```python
num_classes = node.args[1] if len(node.args) > 1 else
node.kwargs.get("num_classes")
if num_classes is None:
raise ValueError("num_classes not found in node.args or node.kwargs")
...
return self.block_builder.emit(relax.op.one_hot(x, on_value, off_value,
num_classes, axis))
```
`num_classes` is a static attribute of `relax.op.one_hot` (it determines
the
output depth), so the frontend is the only place that can report the
problem
usefully. torch itself only rejects an invalid `num_classes` when the
model is
*executed*:
| `num_classes` | native torch | `torch.export` | `fx.symbolic_trace` |
TVM (before) |
|---|---|---|---|---|
| `5` | OK `(3, 5)` | OK | OK | OK `(3, 5)` |
| `0` | `RuntimeError` | **OK** | **OK** | `InternalError: depth must be
positive, but got 0` |
| `-1` (explicit) | OK (infers `max+1`) | OK | OK | `InternalError:
depth must be positive, but got -1` |
| `-2` | `RuntimeError` | OK | OK | `InternalError: depth must be
positive, but got -2` |
`num_classes=-1` is torch's documented "infer the depth from the input"
value.
That cannot be honoured here — the depth would be data dependent and
`relax.op.one_hot`'s depth is static — so it is rejected as well, with
an error
message that says why.
## Fix
Both copies of `_one_hot`
(`python/tvm/relax/frontend/torch/fx_translator.py` and
`python/tvm/relax/frontend/torch/exported_program_translator.py`) gain
the same
check right after the existing "argument missing" guard:
```python
# torch only rejects a non-positive num_classes when the model runs, and
neither
# fx tracing nor export runs it, so the invalid value reaches this
converter.
# num_classes is a static attribute of relax.op.one_hot, so it has to be
rejected
# here rather than by the C++ builder, whose `depth > 0` check never
mentions it.
if isinstance(num_classes, int) and num_classes <= 0:
raise ValueError(
f"one_hot num_classes must be a positive integer, but got
{num_classes}. "
"Inferring the depth from the input (torch's num_classes=-1) is not
"
"supported because the resulting depth is data dependent."
)
```
The `isinstance(num_classes, int)` guard keeps the change conservative:
any
non-literal `num_classes` that may legitimately be dynamic is left
untouched, so
only the reported constant case changes behaviour.
This mirrors the existing frontend-side validation style already used
elsewhere
in the same files (e.g. `_flatten_impl`'s `start_dim`/`end_dim` checks
in
`base_fx_graph_translator.py`), and applies to both entry points, since
`from_fx` uses `fx_translator` and `from_exported_program` uses
`exported_program_translator`.
## Validation
### In-tree regression tests (added)
- `test_one_hot_invalid_num_classes` in
`tests/python/relax/test_frontend_from_fx.py` — `num_classes ∈ {0, -1,
-2}`
are rejected with `ValueError` before lowering; the valid case is
already
covered by the existing `test_one_hot`.
- `test_one_hot_invalid_num_classes` in
`tests/python/relax/test_frontend_from_exported_program.py` —
`num_classes=0` is rejected on the `from_exported_program` path with
`run_ep_decomposition=False`.
Both tests fail without the fix (`tvm.error.InternalError: Check failed:
(depth > 0) is false: one_hot: depth must be positive, but got 0`) and
pass with
it.
Note on the modern path: `from_exported_program` decomposes
`aten.one_hot` to
`arange`/`equal`/`astype` by default (`run_ep_decomposition=True`), so
the
converter is normally bypassed and the test therefore passes
`run_ep_decomposition=False`. Passing that flag is a supported,
already-tested
configuration (see `test_einsum`), and with it `_one_hot` is live code
that hits
exactly the same C++ check.
### Behaviour after the fix
- valid `num_classes ∈ {3, 5, 10}` via `from_fx`: output shape and
values match
native PyTorch exactly (`max|diff| = 0`);
- `num_classes ∈ {0, -1, -2}` via `from_fx`: `ValueError: one_hot
num_classes
must be a positive integer, but got 0. Inferring the depth from the
input
(torch's num_classes=-1) is not supported because the resulting depth is
data
dependent.`
- `num_classes=0` via
`from_exported_program(run_ep_decomposition=False)`: same
`ValueError`.
### Full-suite run
`tests/python/relax/test_frontend_from_fx.py` and
`tests/python/relax/test_frontend_from_exported_program.py` were run in
full
with the change (390 tests: 372 passed, 16 failed, 2 skipped) and again
with the
change reverted (372 passed, 16 failed, 2 skipped, 2 deselected). The
two
failure sets are **identical**, so the change introduces no regressions.
The 16
failures are pre-existing and unrelated to `one_hot`
(`test_extended_unary_ops`,
`test_interpolate`, `test_select_slice`, `test_masked_select`,
`test_to_copy`,
`test_index_put`, `test_eye`, `test_cross_entropy`, the
`test_dynamic_shape*`
family, `test_sym_size_int`, `test_stochastic_depth`); they come from
the local
test tree being newer than the source/lib build used for the run, and
were
verified to fail identically before and after the change.
## Files changed
- `python/tvm/relax/frontend/torch/fx_translator.py` — validate
`num_classes`
in `_one_hot`.
- `python/tvm/relax/frontend/torch/exported_program_translator.py` —
validate
`num_classes` in `_one_hot`.
- `tests/python/relax/test_frontend_from_fx.py` — add
`test_one_hot_invalid_num_classes`.
- `tests/python/relax/test_frontend_from_exported_program.py` — add
`test_one_hot_invalid_num_classes`.
---
.../frontend/torch/exported_program_translator.py | 10 ++++++++++
python/tvm/relax/frontend/torch/fx_translator.py | 10 ++++++++++
.../relax/test_frontend_from_exported_program.py | 15 +++++++++++++++
tests/python/relax/test_frontend_from_fx.py | 20 ++++++++++++++++++++
4 files changed, 55 insertions(+)
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index 18c0ef96d8..f3bfce230a 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1200,6 +1200,16 @@ class ExportedProgramImporter(BaseFXGraphImporter):
num_classes = node.args[1] if len(node.args) > 1 else
node.kwargs.get("num_classes")
if num_classes is None:
raise ValueError("num_classes not found in node.args or
node.kwargs")
+ # torch only rejects a non-positive num_classes when the model runs,
and neither
+ # fx tracing nor export runs it, so the invalid value reaches this
converter.
+ # num_classes is a static attribute of relax.op.one_hot, so it has to
be rejected
+ # here rather than by the C++ builder, whose `depth > 0` check never
mentions it.
+ if isinstance(num_classes, int) and num_classes <= 0:
+ raise ValueError(
+ f"one_hot num_classes must be a positive integer, but got
{num_classes}. "
+ "Inferring the depth from the input (torch's num_classes=-1)
is not "
+ "supported because the resulting depth is data dependent."
+ )
on_value = node.args[2] if len(node.args) > 2 else
node.kwargs.get("on_value", 1)
off_value = node.args[3] if len(node.args) > 3 else
node.kwargs.get("off_value", 0)
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py
b/python/tvm/relax/frontend/torch/fx_translator.py
index 4517650112..2e35ce6ce7 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -727,6 +727,16 @@ class TorchFXImporter(BaseFXGraphImporter):
num_classes = node.args[1] if len(node.args) > 1 else
node.kwargs.get("num_classes")
if num_classes is None:
raise ValueError("num_classes not found in node.args or
node.kwargs")
+ # torch only rejects a non-positive num_classes when the model runs,
and neither
+ # fx tracing nor export runs it, so the invalid value reaches this
converter.
+ # num_classes is a static attribute of relax.op.one_hot, so it has to
be rejected
+ # here rather than by the C++ builder, whose `depth > 0` check never
mentions it.
+ if isinstance(num_classes, int) and num_classes <= 0:
+ raise ValueError(
+ f"one_hot num_classes must be a positive integer, but got
{num_classes}. "
+ "Inferring the depth from the input (torch's num_classes=-1)
is not "
+ "supported because the resulting depth is data dependent."
+ )
on_value = node.args[2] if len(node.args) > 2 else
node.kwargs.get("on_value", 1)
off_value = node.args[3] if len(node.args) > 3 else
node.kwargs.get("off_value", 0)
axis = node.args[4] if len(node.args) > 4 else node.kwargs.get("axis",
-1)
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 43065d9c13..484c68bc66 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -7147,6 +7147,21 @@ def test_one_hot():
verify_model(OneHot(), example_args, {}, Expected)
+def test_one_hot_invalid_num_classes():
+ class OneHot(Module):
+ def forward(self, indices):
+ return torch.nn.functional.one_hot(indices, num_classes=0)
+
+ example_args = (torch.randint(0, 5, (5,), dtype=torch.int64),)
+ exported_program = export(OneHot(), args=example_args)
+
+ # With the default decomposition, one_hot is rewritten to
arange/equal/astype and never
+ # reaches this converter. Without it, the non-positive num_classes must be
rejected by
+ # the frontend instead of failing an internal `depth > 0` check in
relax.op.one_hot.
+ with pytest.raises(ValueError, match="num_classes must be a positive
integer"):
+ from_exported_program(exported_program, run_ep_decomposition=False)
+
+
def test_ones_like():
class OnesLike(Module):
def forward(self, input):
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index 961d863dcc..b7074a5f94 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -5838,6 +5838,26 @@ def test_one_hot():
verify_model(OneHot(), [([5], "int32")], {}, Expected)
+def test_one_hot_invalid_num_classes():
+ input_info = [([5], "int32")]
+
+ class OneHot(Module):
+ def __init__(self, num_classes):
+ super().__init__()
+ self.num_classes = num_classes
+
+ def forward(self, indices):
+ return torch.nn.functional.one_hot(indices,
num_classes=self.num_classes)
+
+ # torch only rejects a non-positive num_classes when the model is
executed, and
+ # fx.symbolic_trace does not execute it, so the invalid value reaches the
frontend and
+ # has to be rejected there instead of failing an internal `depth > 0`
check in
+ # relax.op.one_hot that never mentions num_classes.
+ for num_classes in (0, -1, -2):
+ with pytest.raises(ValueError, match="num_classes must be a positive
integer"):
+ from_fx(fx.symbolic_trace(OneHot(num_classes)), input_info)
+
+
def test_empty_like():
class EmptyLike(Module):
def forward(self, data):
