[Fix][Relax][Frontend][Torch] Validate num_classes in the one_hot converters

[siyiweigeHEW]

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:

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:

# 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.