apache / apache/tvm

[Bug][Relax][Frontend][ONNX] Softmax ignores opset<=12 coerce-to-2D semantics and defaults `axis` to -1 instead of 1, silently producing wrong output for valid models

Open
#20,183 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

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

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in python/tvm/relax/frontend/onnx/onnx_frontend.py at the Softmax converter around lines 613-615 and the version selection logic around lines 288-308. Run the provided opset 11 reproduction and compare TVM with onnxruntime for the listed axis cases. Done means Softmax, LogSoftmax, and Hardmax match the version-specific ONNX results for opsets ≤12 while preserving opset ≥13 behavior.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.