siyiweigeHEW opened a new issue, #20183:
URL: https://github.com/apache/tvm/issues/20183
### Expected behavior
A valid ONNX `Softmax` node must follow the version-specific semantics:
- **opset ≤ 12** (`Softmax-1` / `Softmax-11`): the input is coerced into a
2-D tensor of shape `[a_0 * ... * a_{k-1}, a_k * ... * a_{n-1}]` (where `k` is
`axis`), softmax is computed along the second dimension of the coerced tensor,
and the result is reshaped back to the original shape. The default `axis` is
**1**.
- **opset ≥ 13** (`Softmax-13`): softmax is computed directly along `axis`;
the default `axis` is **-1**.
For example, with input `X` of shape `(2, 3, 4)` at opset 11 and `axis = 1`,
the output must be `exp(x) / sum(exp(x), axis=(1, 2), keepdims=True)` (12
elements normalized per row), because dims `1, 2` are flattened into the
coerced second dimension. The model passes `onnx.checker` and runs correctly in
onnxruntime.
### Actual behavior
`tvm.relax.frontend.onnx.from_onnx` silently returns numerically **wrong**
output for any opset ≤ 12 model whose input has rank > 2 with a non-last axis
(or no `axis` attribute).
The `Softmax` converter only registers `_impl_v13`
([onnx_frontend.py:613-615](https://github.com/apache/tvm/blob/262c6d2e04/python/tvm/relax/frontend/onnx/onnx_frontend.py#L613-L615)):
```python
class Softmax(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
axis = attr.get("axis", -1)
return relax.op.nn.softmax(inputs[0], axis=axis)
```
Since `OnnxOpConverter.get_converter`
([onnx_frontend.py:288-308](https://github.com/apache/tvm/blob/262c6d2e04/python/tvm/relax/frontend/onnx/onnx_frontend.py#L288-L308))
picks the largest registered `_impl_v*` ≤ the model opset, *every* opset (1,
11, 12, …) is routed to `_impl_v13`. That implementation:
1. never performs the required coerce-to-2D — it runs direct-axis softmax
over the single `axis` dimension, and
2. defaults `axis` to **-1** instead of the spec default **1** for opset ≤
12 (confirmed by `onnx.defs`: `Softmax-1`/`Softmax-11` default `axis=1`,
`Softmax-13` default `axis=-1`).
As a result valid opset ≤ 12 models are **silently** (no error) computed
with the wrong normalization. For a `(2,3,4)` input at opset 11:
| case | onnxruntime (spec: coerce-to-2D, default axis=1) | TVM relax
(direct axis) | max \|diff\| |
|------|---------------------------------------------------|--------------------------|---------------|
| no `axis` | softmax over coerced `[2,12]` (12 elems) | softmax over last 4
elems | 0.644 |
| `axis=1` | softmax over coerced `[2,12]` (12 elems) | softmax over
`axis=1` (3 elems) | 0.950 |
| `axis=0` | softmax over coerced `[1,24]` (24 elems) | softmax over
`axis=0` (2 elems) | 1.000 |
`LogSoftmax` and `Hardmax` are affected identically (same `_impl_v13`-only
pattern).
This was already recognized upstream: apache/tvm PR #19428
("[Relax][FRONTEND][ONNX] Support Softmax, LogSoftmax and Hardmax when opset
version ≤12", commit `7eea6df1b6`) fixes it by adding `_impl_v1`/`_impl_v11`
with flatten-to-2D + reshape-back. The version below (built 2026-02-11)
predates that fix.
### Environment
- OS: Linux
- TVM: v0.24.dev0 (main branch, commit `262c6d2e0`, built 2026-02-11)
- Python: 3.11
- onnx: 1.20.1
- onnxruntime: 1.24.1
### Steps to reproduce
```python
"""Repro: ONNX Softmax opset<=12 coerce-to-2D semantics silently wrong in
TVM relax frontend."""
import numpy as np
import onnx, onnxruntime
from onnx import helper, TensorProto
import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
def build(shape, axis, opset):
node = helper.make_node("Softmax", ["X"], ["Y"])
if axis is not None:
node.attribute.append(helper.make_attribute("axis", axis))
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, list(shape))
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)
g = helper.make_graph([node], "g", [X], [Y])
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", opset)])
m.ir_version = 8
return m
def run_tvm(model, shape):
mod = from_onnx(model, shape_dict={"X": list(shape)})
ex = relax.build(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
return vm["main"](x).numpy()
def show(label, shape, axis, opset):
m = build(shape, axis, opset)
o = onnxruntime.InferenceSession(m.SerializeToString(),
providers=["CPUExecutionProvider"]).run(None, {"X": x})[0]
t = run_tvm(m, shape)
d = float(np.abs(t.astype(np.float64) - o.astype(np.float64)).max())
print(f"{label} onnxruntime shape={list(o.shape)}
max|diff|(tvm-vs-ort)={d:.3e}")
x = (np.arange(24).reshape(2, 3, 4) + 1).astype(np.float32)
show("Case1 (op11, no axis)", (2, 3, 4), None, 11) # ONNX default axis=1
show("Case2 (op11, axis=1) ", (2, 3, 4), 1, 11)
```
Actual output:
```
Case1 (op11, no axis) onnxruntime shape=[2, 3, 4]
max|diff|(tvm-vs-ort)=6.437e-01
Case2 (op11, axis=1) onnxruntime shape=[2, 3, 4]
max|diff|(tvm-vs-ort)=9.502e-01
```
### Triage
* needs-triage
* bug
* relax
* frontend/onnx
--
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]
