Repository navigation
[Bug] layer_norm should support broadcast in all axes #20802
Description
https://onnx.ai/onnx/operators/onnx__LayerNormalization.html allows tensors Scale and B should be unidirectional broadcastable to tensor X, but layer_norm op in tvm currently require input to have size same as the lengths of the data on the given axes (see https://github.com/apache/tvm/blob/main/src/relax/op/nn/nn.cc#L666). This results in error if a model using such broadcast rule is imported.
Example reproduce:
import tvm, onnx
from tvm.relax.frontend.onnx import from_onnx
model = onnx.load_model_from_string(b'\x08\x08\x12\x0cbackend-test:\xd4\x01\nS\n\x01X\n\x01W\n\x01B\x12\x01Y\x12\x04Mean\x12\tInvStdDev\x1a\x13LayerNormalization0"\x12LayerNormalization*\x0b\n\x04axis\x18\x01\xa0\x01\x02\x12!test_layer_normalization_4d_axis1Z\x17\n\x01X\x12\x12\n\x10\x08\x01\x12\x0c\n\x02\x08\x02\n\x02\x08\x03\n\x02\x08\x05Z\x13\n\x01W\x12\x0e\n\x0c\x08\x01\x12\x08\n\x02\x08\x03\n\x02\x08\x01Z\x13\n\x01B\x12\x0e\n\x0c\x08\x01\x12\x08\n\x02\x08\x03\n\x02\x08\x01b\x17\n\x01Y\x12\x12\n\x10\x08\x01\x12\x0c\n\x02\x08\x02\n\x02\x08\x03\n\x02\x08\x05B\x04\n\x00\x10\x11')
tvm.relax.transform.ConvertLayout({})(from_onnx(model))
Activity
xgbtxy commentedon Oct 20, 2026
I reproduced the importer failure with the Windows x64
apache-tvm==0.27.0.post1wheel and ONNX 1.23.2, using explicitly constructed models rather than the serialized example.For
X=[2,3,5],axis=1, the[3,5]Scale/B control imports, while[3,1]Scale/B is rejected by the Relaxlayer_normparameter-shape check. Separate cases with a trailing vector, scalar affine parameters, independently broadcastable Scale/B, batch-dependent parameters, omitted bias, andaxis=-2are also rejected. All eight models pass ONNX full checking and strict shape inference, and ONNX ReferenceEvaluator matches a NumPy oracle. This checks TVM import behavior; I have not compiled/executed these models with TVM.The current frontend also requires Scale/B to have identical shapes, although ONNX allows each to broadcast independently. A frontend-only follow-up could preserve the existing Relax operator contract; broadcasting only to the normalized suffix would fix the reported shape but would not cover batch-dependent affine parameters.
I noticed #20510 is already updating this converter's normalization axes and optional outputs. I have not started a competing implementation. Would a separate broadcast follow-up after that PR be useful, or is this intended to be included in its scope?
Disclosure: this reproduction and source analysis were performed with AI assistance (OpenAI Codex).