siyiweigeHEW opened a new pull request, #20320:
URL: https://github.com/apache/tvm/pull/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`.
--
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]
