Skip to content

[FIX][Relax][ONNX] Correct normalization semantics - #20510

Open
Yuhx141 wants to merge 1 commit into
apache:mainfrom
Yuhx141:yuhx/fix-onnx-normalization-20261001
Open

Yuhx141 wants to merge 1 commit into
apache:mainfrom
Yuhx141:yuhx/fix-onnx-normalization-20261001

Conversation

@Yuhx141

@Yuhx141 Yuhx141 commented Oct 1, 2026

Copy link
Copy Markdown

This fixes three normalization correctness issues at their respective contract boundaries:

  • convert ONNX BatchNormalization.momentum to Relax's update-coefficient convention;
  • return the real ONNX LayerNormalization Mean and InvStdDev outputs instead of scalar placeholders;
  • normalize a negative BatchNorm axis before DecomposeOpsForTraining constructs reduction axes.

The LayerNormalization statistics follow ONNX stash_type=1 and are computed in float32. The BatchNorm momentum conversion is kept in the ONNX frontend so existing Relax/TOPI semantics remain unchanged.

Closes #20496
Closes #20503
Closes #20504

Tests:

  • pytest tests/python/relax/test_frontend_onnx.py -k "layer_norm or batch_norm" -q (8 passed)
  • pytest tests/python/relax/test_transform_decompose_ops.py::test_batch_norm_training_negative_axis -q (1 passed)

@Yuhx141

Yuhx141 commented Oct 15, 2026

Copy link
Copy Markdown
Author

Hi @tlopex, could you take a look at these Relax/ONNX normalization correctness fixes when you have a chance? Please let me know if any further changes are needed. Thanks!

@Yuhx141

Yuhx141 commented Oct 31, 2026

Copy link
Copy Markdown
Author

@tqchen @mshr-h

Could one of you take a look at these Relax/ONNX normalization correctness fixes when you have a chance, or point me to the right reviewer? I am keen to help get these fixes into TVM and happy to address feedback, add further regression tests, rebase, or split the changes if that would make review easier.

The GitHub Actions checks also seem to need maintainer help with fork-workflow approval or a restart. I would appreciate any guidance on getting the appropriate checks running. Thank you for your time!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant